From 3d86e5f09714d654c81df6670b3c58e110877e42 Mon Sep 17 00:00:00 2001 From: Yuki Watanabe <31463517+B-Step62@users.noreply.github.com> Date: Wed, 22 May 2024 21:53:24 +0900 Subject: [PATCH] Fix cascading behavior for trace tag table (#12102) Signed-off-by: B-Step62 --- ...add_cascade_deletion_to_trace_tables_fk.py | 37 +++++++++++++++++++ mlflow/store/tracking/dbmodels/models.py | 8 +++- tests/db/schemas/mssql.sql | 2 +- tests/db/schemas/mysql.sql | 4 +- tests/db/schemas/postgresql.sql | 4 +- tests/db/schemas/sqlite.sql | 4 +- tests/resources/db/latest_schema.sql | 4 +- tests/store/tracking/test_sqlalchemy_store.py | 8 +++- 8 files changed, 58 insertions(+), 13 deletions(-) create mode 100644 mlflow/store/db_migrations/versions/5b0e9adcef9c_add_cascade_deletion_to_trace_tables_fk.py diff --git a/mlflow/store/db_migrations/versions/5b0e9adcef9c_add_cascade_deletion_to_trace_tables_fk.py b/mlflow/store/db_migrations/versions/5b0e9adcef9c_add_cascade_deletion_to_trace_tables_fk.py new file mode 100644 index 0000000000000..bb1ea08015f10 --- /dev/null +++ b/mlflow/store/db_migrations/versions/5b0e9adcef9c_add_cascade_deletion_to_trace_tables_fk.py @@ -0,0 +1,37 @@ +"""add cascade deletion to trace tables foreign keys + +Revision ID: 5b0e9adcef9c +Revises: 867495a8f9d4 +Create Date: 2024-05-22 17:44:24.597019 + +""" +from alembic import op +from mlflow.store.tracking.dbmodels.models import SqlTraceInfo, SqlTraceRequestMetadata, SqlTraceTag + + +# revision identifiers, used by Alembic. +revision = '5b0e9adcef9c' +down_revision = '867495a8f9d4' +branch_labels = None +depends_on = None + + +def upgrade(): + tables = [SqlTraceTag.__tablename__, SqlTraceRequestMetadata.__tablename__] + for table in tables: + fk_tag_constaint_name = f"fk_{table}_request_id" + # We have to use batch_alter_table as SQLite does not support ALTER outside of a batch operation. + with op.batch_alter_table(table, schema=None) as batch_op: + batch_op.drop_constraint(fk_tag_constaint_name, type_="foreignkey") + batch_op.create_foreign_key( + fk_tag_constaint_name, + SqlTraceInfo.__tablename__, + ["request_id"], + ["request_id"], + # Add cascade deletion to the foreign key constraint. This is the only change in this migration. + ondelete="CASCADE", + ) + + +def downgrade(): + pass diff --git a/mlflow/store/tracking/dbmodels/models.py b/mlflow/store/tracking/dbmodels/models.py index 0794c7d32bdfe..2378e1a14d01b 100644 --- a/mlflow/store/tracking/dbmodels/models.py +++ b/mlflow/store/tracking/dbmodels/models.py @@ -706,7 +706,9 @@ class SqlTraceTag(Base): """ Value associated with tag: `String` (limit 250 characters). Could be *null*. """ - request_id = Column(String(50), ForeignKey("trace_info.request_id"), nullable=False) + request_id = Column( + String(50), ForeignKey("trace_info.request_id", ondelete="CASCADE"), nullable=False + ) """ Request ID to which this tag belongs: *Foreign Key* into ``trace_info`` table. """ @@ -734,7 +736,9 @@ class SqlTraceRequestMetadata(Base): """ Value associated with metadata: `String` (limit 250 characters). Could be *null*. """ - request_id = Column(String(50), ForeignKey("trace_info.request_id"), nullable=False) + request_id = Column( + String(50), ForeignKey("trace_info.request_id", ondelete="CASCADE"), nullable=False + ) """ Request ID to which this metadata belongs: *Foreign Key* into ``trace_info`` table. """ diff --git a/tests/db/schemas/mssql.sql b/tests/db/schemas/mssql.sql index 171f86b779b19..428983f399922 100644 --- a/tests/db/schemas/mssql.sql +++ b/tests/db/schemas/mssql.sql @@ -191,7 +191,7 @@ CREATE TABLE trace_request_metadata ( value VARCHAR(8000) COLLATE "SQL_Latin1_General_CP1_CI_AS", request_id VARCHAR(50) COLLATE "SQL_Latin1_General_CP1_CI_AS" NOT NULL, CONSTRAINT trace_request_metadata_pk PRIMARY KEY (key, request_id), - CONSTRAINT fk_trace_request_metadata_request_id FOREIGN KEY(request_id) REFERENCES trace_info (request_id) + CONSTRAINT fk_trace_request_metadata_request_id FOREIGN KEY(request_id) REFERENCES trace_info (request_id) ON DELETE CASCADE ) diff --git a/tests/db/schemas/mysql.sql b/tests/db/schemas/mysql.sql index c5d539c9d6e20..945c584ed75f3 100644 --- a/tests/db/schemas/mysql.sql +++ b/tests/db/schemas/mysql.sql @@ -198,7 +198,7 @@ CREATE TABLE trace_request_metadata ( value VARCHAR(8000), request_id VARCHAR(50) NOT NULL, CONSTRAINT trace_request_metadata_pk PRIMARY KEY (key, request_id), - CONSTRAINT fk_trace_request_metadata_request_id FOREIGN KEY(request_id) REFERENCES trace_info (request_id) + CONSTRAINT fk_trace_request_metadata_request_id FOREIGN KEY(request_id) REFERENCES trace_info (request_id) ON DELETE CASCADE ) @@ -207,5 +207,5 @@ CREATE TABLE trace_tags ( value VARCHAR(8000), request_id VARCHAR(50) NOT NULL, CONSTRAINT trace_tag_pk PRIMARY KEY (key, request_id), - CONSTRAINT fk_trace_tags_request_id FOREIGN KEY(request_id) REFERENCES trace_info (request_id) + CONSTRAINT fk_trace_tags_request_id FOREIGN KEY(request_id) REFERENCES trace_info (request_id) ON DELETE CASCADE ) diff --git a/tests/db/schemas/postgresql.sql b/tests/db/schemas/postgresql.sql index 99aaf19340e9f..82310ddc2a71d 100644 --- a/tests/db/schemas/postgresql.sql +++ b/tests/db/schemas/postgresql.sql @@ -196,7 +196,7 @@ CREATE TABLE trace_request_metadata ( value VARCHAR(8000), request_id VARCHAR(50) NOT NULL, CONSTRAINT trace_request_metadata_pk PRIMARY KEY (key, request_id), - CONSTRAINT fk_trace_request_metadata_request_id FOREIGN KEY(request_id) REFERENCES trace_info (request_id) + CONSTRAINT fk_trace_request_metadata_request_id FOREIGN KEY(request_id) REFERENCES trace_info (request_id) ON DELETE CASCADE ) @@ -205,5 +205,5 @@ CREATE TABLE trace_tags ( value VARCHAR(8000), request_id VARCHAR(50) NOT NULL, CONSTRAINT trace_tag_pk PRIMARY KEY (key, request_id), - CONSTRAINT fk_trace_tags_request_id FOREIGN KEY(request_id) REFERENCES trace_info (request_id) + CONSTRAINT fk_trace_tags_request_id FOREIGN KEY(request_id) REFERENCES trace_info (request_id) ON DELETE CASCADE ) diff --git a/tests/db/schemas/sqlite.sql b/tests/db/schemas/sqlite.sql index e7eda13ad4dfb..a690cac192943 100644 --- a/tests/db/schemas/sqlite.sql +++ b/tests/db/schemas/sqlite.sql @@ -199,7 +199,7 @@ CREATE TABLE trace_request_metadata ( value VARCHAR(8000), request_id VARCHAR(50) NOT NULL, CONSTRAINT trace_request_metadata_pk PRIMARY KEY (key, request_id), - CONSTRAINT fk_trace_request_metadata_request_id FOREIGN KEY(request_id) REFERENCES trace_info (request_id) + CONSTRAINT fk_trace_request_metadata_request_id FOREIGN KEY(request_id) REFERENCES trace_info (request_id) ON DELETE CASCADE ) @@ -208,5 +208,5 @@ CREATE TABLE trace_tags ( value VARCHAR(8000), request_id VARCHAR(50) NOT NULL, CONSTRAINT trace_tag_pk PRIMARY KEY (key, request_id), - CONSTRAINT fk_trace_tags_request_id FOREIGN KEY(request_id) REFERENCES trace_info (request_id) + CONSTRAINT fk_trace_tags_request_id FOREIGN KEY(request_id) REFERENCES trace_info (request_id) ON DELETE CASCADE ) diff --git a/tests/resources/db/latest_schema.sql b/tests/resources/db/latest_schema.sql index 242ea1e0582a4..353cd544e5aec 100644 --- a/tests/resources/db/latest_schema.sql +++ b/tests/resources/db/latest_schema.sql @@ -199,7 +199,7 @@ CREATE TABLE trace_request_metadata ( value VARCHAR(8000), request_id VARCHAR(50) NOT NULL, CONSTRAINT trace_request_metadata_pk PRIMARY KEY (key, request_id), - CONSTRAINT fk_trace_request_metadata_request_id FOREIGN KEY(request_id) REFERENCES trace_info (request_id) + CONSTRAINT fk_trace_request_metadata_request_id FOREIGN KEY(request_id) REFERENCES trace_info (request_id) ON DELETE CASCADE ) @@ -208,6 +208,6 @@ CREATE TABLE trace_tags ( value VARCHAR(8000), request_id VARCHAR(50) NOT NULL, CONSTRAINT trace_tag_pk PRIMARY KEY (key, request_id), - CONSTRAINT fk_trace_tags_request_id FOREIGN KEY(request_id) REFERENCES trace_info (request_id) + CONSTRAINT fk_trace_tags_request_id FOREIGN KEY(request_id) REFERENCES trace_info (request_id) ON DELETE CASCADE ) diff --git a/tests/store/tracking/test_sqlalchemy_store.py b/tests/store/tracking/test_sqlalchemy_store.py index bf94954a46f13..48f482688d12c 100644 --- a/tests/store/tracking/test_sqlalchemy_store.py +++ b/tests/store/tracking/test_sqlalchemy_store.py @@ -4148,8 +4148,12 @@ def test_delete_traces(store): exp2 = store.create_experiment("exp2") for i in range(10): - _create_trace(store, f"tr-exp1-{i}", exp1) - _create_trace(store, f"tr-exp2-{i}", exp2) + _create_trace( + store, f"tr-exp1-{i}", exp1, tags={"tag": "apple"}, request_metadata={"rq": "foo"} + ) + _create_trace( + store, f"tr-exp2-{i}", exp2, tags={"tag": "orange"}, request_metadata={"rq": "bar"} + ) traces, _ = store.search_traces([exp1, exp2]) assert len(traces) == 20