chore(datasets): Refactor DatasetDAO update to leverage bulk logic for create, update, and delete operations (#22957)

This commit is contained in:
John Bodley 2023-02-14 06:47:56 +13:00 committed by GitHub
parent 4980621902
commit f3bdcdc79b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -187,39 +187,50 @@ class DatasetDAO(BaseDAO): # pylint: disable=too-many-public-methods
then we delete. then we delete.
""" """
column_by_id = {column.id: column for column in model.columns}
seen = set()
original_cols = {obj.id for obj in model.columns}
if override_columns: if override_columns:
for id_ in original_cols: db.session.query(TableColumn).filter(
DatasetDAO.delete_column(column_by_id[id_], commit=False) TableColumn.table_id == model.id
).delete(synchronize_session="fetch")
db.session.flush() db.session.bulk_insert_mappings(
TableColumn,
for properties in property_columns: [
DatasetDAO.create_column( {**properties, "table_id": model.id}
{**properties, "table_id": model.id}, for properties in property_columns
commit=False, ],
) )
else: else:
for properties in property_columns: columns_by_id = {column.id: column for column in model.columns}
if "id" in properties:
seen.add(properties["id"])
DatasetDAO.update_column( property_columns_by_id = {
column_by_id[properties["id"]], properties["id"]: properties
properties, for properties in property_columns
commit=False, if "id" in properties
) }
else:
DatasetDAO.create_column(
{**properties, "table_id": model.id},
commit=False,
)
for id_ in {obj.id for obj in model.columns} - seen: db.session.bulk_insert_mappings(
DatasetDAO.delete_column(column_by_id[id_], commit=False) TableColumn,
[
{**properties, "table_id": model.id}
for properties in property_columns
if not "id" in properties
],
)
db.session.bulk_update_mappings(
TableColumn,
[
{**columns_by_id[properties["id"]].__dict__, **properties}
for properties in property_columns_by_id.values()
],
)
db.session.query(TableColumn).filter(
TableColumn.id.in_(
{column.id for column in model.columns}
- property_columns_by_id.keys()
)
).delete(synchronize_session="fetch")
if commit: if commit:
db.session.commit() db.session.commit()
@ -241,26 +252,36 @@ class DatasetDAO(BaseDAO): # pylint: disable=too-many-public-methods
then we delete. then we delete.
""" """
metric_by_id = {metric.id: metric for metric in model.metrics} metrics_by_id = {metric.id: metric for metric in model.metrics}
seen = set()
for properties in property_metrics: property_metrics_by_id = {
if "id" in properties: properties["id"]: properties
seen.add(properties["id"]) for properties in property_metrics
if "id" in properties
}
DatasetDAO.update_metric( db.session.bulk_insert_mappings(
metric_by_id[properties["id"]], SqlMetric,
properties, [
commit=False, {**properties, "table_id": model.id}
) for properties in property_metrics
else: if not "id" in properties
DatasetDAO.create_metric( ],
{**properties, "table_id": model.id}, )
commit=False,
)
for id_ in {obj.id for obj in model.metrics} - seen: db.session.bulk_update_mappings(
DatasetDAO.delete_column(metric_by_id[id_], commit=False) SqlMetric,
[
{**metrics_by_id[properties["id"]].__dict__, **properties}
for properties in property_metrics_by_id.values()
],
)
db.session.query(SqlMetric).filter(
SqlMetric.id.in_(
{metric.id for metric in model.metrics} - property_metrics_by_id.keys()
)
).delete(synchronize_session="fetch")
if commit: if commit:
db.session.commit() db.session.commit()