Repository: mlflow/mlflow

Issue:
[BUG] mlflow gc do not remove deleted experiments when runs were tracked with datasets
### Issues Policy acknowledgement

- [X] I have read and agree to submit bug reports in accordance with the [issues policy](https://www.github.com/mlflow/mlflow/blob/master/ISSUE_POLICY.md)

### Where did you encounter this bug?

Local machine

### Willingness to contribute

Yes. I would be willing to contribute a fix for this bug with guidance from the MLflow community.

### MLflow version

- Client: 2.9.2
- Tracking server: 2.9.2


### System information

- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**: Ubuntu 22.04.3 LTS 
- **Python version**: Python 3.10.12
- **yarn version, if running the dev UI**: -


### Describe the problem

"mlflow gc" command do not remove deleted experiments when runs were tracked with datasets

### Tracking information

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
It was a default setup from "Remote Experiment Tracking with MLflow Tracking Server" scenario:
https://mlflow.org/docs/latest/tracking/tutorials/remote-server.html

1. When created, I have run a sample code (attached in below section - "Code to reproduce issue") to generate a "test" experiment.
2. Then, I have deleted it with the web UI.
3. After that I have run below commands:
```shell
export MLFLOW_TRACKING_URI="http://localhost:5000"
mlflow gc --backend-store-uri="postgresql://user:password@localhost:5432/mlflowdb"
```


### Code to reproduce issue

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```python
import mlflow

from sklearn.model_selection import train_test_split
from sklearn.datasets import load_diabetes
from sklearn.ensemble import RandomForestRegressor

mlflow.set_tracking_uri("http://localhost:5000")
mlflow.set_experiment("test")
mlflow.autolog()

db = load_diabetes()
X_train, X_test, y_train, y_test = train_test_split(db.data, db.target)

# Create and train models.
rf = RandomForestRegressor(n_estimators=100, max_depth=6, max_features=3)
rf.fit(X_train, y_train)

# Use the model to make predictions on the test dataset.
predictions = rf.predict(X_test)

```


### Stack trace

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```
mlflow gc --backend-store-uri="postgresql://user:password@localhost:5432/mlflowdb"
Run with ID 6fb6bb680a2249388ec8346febc6540c has been permanently deleted.
Traceback (most recent call last):
  File "/home/mlokos/dev/mlflow/venv/lib/python3.10/site-packages/sqlalchemy/engine/base.py", line 1969, in _exec_single_context
    self.dialect.do_execute(
  File "/home/mlokos/dev/mlflow/venv/lib/python3.10/site-packages/sqlalchemy/engine/default.py", line 922, in do_execute
    cursor.execute(statement, parameters)
psycopg2.errors.ForeignKeyViolation: update or delete on table "experiments" violates foreign key constraint "datasets_experiment_id_fkey" on table "datasets"
DETAIL:  Key (experiment_id)=(1) is still referenced from table "datasets".


The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/mlokos/dev/mlflow/venv/lib/python3.10/site-packages/mlflow/store/db/utils.py", line 142, in make_managed_session
    session.commit()
  File "/home/mlokos/dev/mlflow/venv/lib/python3.10/site-packages/sqlalchemy/orm/session.py", line 1969, in commit
    trans.commit(_to_root=True)
  File "<string>", line 2, in commit
  File "/home/mlokos/dev/mlflow/venv/lib/python3.10/site-packages/sqlalchemy/orm/state_changes.py", line 139, in _go
    ret_value = fn(self, *arg, **kw)
  File "/home/mlokos/dev/mlflow/venv/lib/python3.10/site-packages/sqlalchemy/orm/session.py", line 1256, in commit
    self._prepare_impl()
  File "<string>", line 2, in _prepare_impl
  File "/home/mlokos/dev/mlflow/venv/lib/python3.10/site-packages/sqlalchemy/orm/state_changes.py", line 139, in _go
    ret_value = fn(self, *arg, **kw)
  File "/home/mlokos/dev/mlflow/venv/lib/python3.10/site-packages/sqlalchemy/orm/session.py", line 1231, in _prepare_impl
    self.session.flush()
  File "/home/mlokos/dev/mlflow/venv/lib/python3.10/site-packages/sqlalchemy/orm/session.py", line 4312, in flush
    self._flush(objects)
  File "/home/mlokos/dev/mlflow/venv/lib/python3.10/site-packages/sqlalchemy/orm/session.py", line 4447, in _flush
    with util.safe_reraise():
  File "/home/mlokos/dev/mlflow/venv/lib/python3.10/site-packages/sqlalchemy/util/langhelpers.py", line 146, in __exit__
    raise exc_value.with_traceback(exc_tb)
  File "/home/mlokos/dev/mlflow/venv/lib/python3.10/site-packages/sqlalchemy/orm/session.py", line 4408, in _flush
    flush_context.execute()
  File "/home/mlokos/dev/mlflow/venv/lib/python3.10/site-packages/sqlalchemy/orm/unitofwork.py", line 466, in execute
    rec.execute(self)
  File "/home/mlokos/dev/mlflow/venv/lib/python3.10/site-packages/sqlalchemy/orm/unitofwork.py", line 679, in execute
    util.preloaded.orm_persistence.delete_obj(
  File "/home/mlokos/dev/mlflow/venv/lib/python3.10/site-packages/sqlalchemy/orm/persistence.py", line 191, in delete_obj
    _emit_delete_statements(
  File "/home/mlokos/dev/mlflow/venv/lib/python3.10/site-packages/sqlalchemy/orm/persistence.py", line 1456, in _emit_delete_statements
    c = connection.execute(
  File "/home/mlokos/dev/mlflow/venv/lib/python3.10/site-packages/sqlalchemy/engine/base.py", line 1416, in execute
    return meth(
  File "/home/mlokos/dev/mlflow/venv/lib/python3.10/site-packages/sqlalchemy/sql/elements.py", line 516, in _execute_on_connection
    return connection._execute_clauseelement(
  File "/home/mlokos/dev/mlflow/venv/lib/python3.10/site-packages/sqlalchemy/engine/base.py", line 1639, in _execute_clauseelement
    ret = self._execute_context(
  File "/home/mlokos/dev/mlflow/venv/lib/python3.10/site-packages/sqlalchemy/engine/base.py", line 1848, in _execute_context
    return self._exec_single_context(
  File "/home/mlokos/dev/mlflow/venv/lib/python3.10/site-packages/sqlalchemy/engine/base.py", line 1988, in _exec_single_context
    self._handle_dbapi_exception(
  File "/home/mlokos/dev/mlflow/venv/lib/python3.10/site-packages/sqlalchemy/engine/base.py", line 2343, in _handle_dbapi_exception
    raise sqlalchemy_exception.with_traceback(exc_info[2]) from e
  File "/home/mlokos/dev/mlflow/venv/lib/python3.10/site-packages/sqlalchemy/engine/base.py", line 1969, in _exec_single_context
    self.dialect.do_execute(
  File "/home/mlokos/dev/mlflow/venv/lib/python3.10/site-packages/sqlalchemy/engine/default.py", line 922, in do_execute
    cursor.execute(statement, parameters)
sqlalchemy.exc.IntegrityError: (psycopg2.errors.ForeignKeyViolation) update or delete on table "experiments" violates foreign key constraint "datasets_experiment_id_fkey" on table "datasets"
DETAIL:  Key (experiment_id)=(1) is still referenced from table "datasets".

[SQL: DELETE FROM experiments WHERE experiments.experiment_id = %(experiment_id)s]
[parameters: {'experiment_id': 1}]
(Background on this error at: https://sqlalche.me/e/20/gkpj)

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/mlokos/dev/mlflow/venv/bin/mlflow", line 8, in <module>
    sys.exit(cli())
  File "/home/mlokos/dev/mlflow/venv/lib/python3.10/site-packages/click/core.py", line 1157, in __call__
    return self.main(*args, **kwargs)
  File "/home/mlokos/dev/mlflow/venv/lib/python3.10/site-packages/click/core.py", line 1078, in main
    rv = self.invoke(ctx)
  File "/home/mlokos/dev/mlflow/venv/lib/python3.10/site-packages/click/core.py", line 1688, in invoke
    return _process_result(sub_ctx.command.invoke(sub_ctx))
  File "/home/mlokos/dev/mlflow/venv/lib/python3.10/site-packages/click/core.py", line 1434, in invoke
    return ctx.invoke(self.callback, **ctx.params)
  File "/home/mlokos/dev/mlflow/venv/lib/python3.10/site-packages/click/core.py", line 783, in invoke
    return __callback(*args, **kwargs)
  File "/home/mlokos/dev/mlflow/venv/lib/python3.10/site-packages/mlflow/cli.py", line 630, in gc
    backend_store._hard_delete_experiment(experiment_id)
  File "/home/mlokos/dev/mlflow/venv/lib/python3.10/site-packages/mlflow/store/tracking/sqlalchemy_store.py", line 421, in _hard_delete_experiment
    with self.ManagedSessionMaker() as session:
  File "/usr/lib/python3.10/contextlib.py", line 142, in __exit__
    next(self.gen)
  File "/home/mlokos/dev/mlflow/venv/lib/python3.10/site-packages/mlflow/store/db/utils.py", line 155, in make_managed_session
    raise MlflowException(message=e, error_code=BAD_REQUEST)
mlflow.exceptions.MlflowException: (psycopg2.errors.ForeignKeyViolation) update or delete on table "experiments" violates foreign key constraint "datasets_experiment_id_fkey" on table "datasets"
DETAIL:  Key (experiment_id)=(1) is still referenced from table "datasets".

[SQL: DELETE FROM experiments WHERE experiments.experiment_id = %(experiment_id)s]
[parameters: {'experiment_id': 1}]
(Background on this error at: https://sqlalche.me/e/20/gkpj)

```


### Other info / logs

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
The issue is solely connected with datasets-table-handling.
I was able to make "mlflow gc" command work, but had to manually remove rows form that table.
```shell
psql -d mlflowdb -U user --password -h localhost -p 5432
truncate TABLE datasets; # inside psql cli
```
After that the "mlflow gc" command was able to remove data from PostgreSQL and artifacts from minIO.
```shell
export MLFLOW_TRACKING_URI="http://localhost:5000"
mlflow gc --backend-store-uri="postgresql://user:password@localhost:5432/mlflowdb"
Experiment with ID 1 has been permanently deleted.
```

### What component(s) does this bug affect?

- [ ] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/build`: Build and test infrastructure for MLflow
- [X] `area/deployments`: MLflow Deployments client APIs, server, and third-party Deployments integrations
- [ ] `area/docs`: MLflow documentation pages
- [X] `area/examples`: Example code
- [ ] `area/model-registry`: Model Registry service, APIs, and the fluent client calls for Model Registry
- [ ] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/recipes`: Recipes, Recipe APIs, Recipe configs, Recipe Templates
- [ ] `area/projects`: MLproject format, project running backends
- [ ] `area/scoring`: MLflow Model server, model deployment tools, Spark UDFs
- [X] `area/server-infra`: MLflow Tracking server backend
- [X] `area/tracking`: Tracking Service, tracking client APIs, autologging

### What interface(s) does this bug affect?

- [ ] `area/uiux`: Front-end, user experience, plotting, JavaScript, JavaScript dev server
- [ ] `area/docker`: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
- [X] `area/sqlalchemy`: Use of SQLAlchemy in the Tracking Service or Model Registry
- [ ] `area/windows`: Windows support

### What language(s) does this bug affect?

- [ ] `language/r`: R APIs and clients
- [ ] `language/java`: Java APIs and clients
- [ ] `language/new`: Proposals for new client languages

### What integration(s) does this bug affect?

- [ ] `integrations/azure`: Azure and Azure ML integrations
- [ ] `integrations/sagemaker`: SageMaker integrations
- [ ] `integrations/databricks`: Databricks integrations


Discussion:
Thanks for the report! I can reproduce this error, will investigate and get back to you!
Hello folks. Same bug was reproducted today, on the same remote tracking server setup, but with a MySQL database and GCS bucket storage.

A possible workaround I implemented on my system this morning is to flush partially the "datasets" table, before calling `mlflow gc`, with the following SQL command:
`DELETE datasets FROM datasets JOIN experiments ON datasets.experiment_id = experiments.experiment_id WHERE experiments.lifecycle_stage = "deleted";`
<!-- assign-maintainer -->
@mlflow/mlflow-team Please assign a maintainer and start triaging this issue.

Resolving patch:
diff --git a/mlflow/store/db_migrations/versions/93bd306463bd_add_cascading_delete_on_experiments_fk_in_datasets_table.py b/mlflow/store/db_migrations/versions/93bd306463bd_add_cascading_delete_on_experiments_fk_in_datasets_table.py
new file mode 100644
--- /dev/null
+++ b/mlflow/store/db_migrations/versions/93bd306463bd_add_cascading_delete_on_experiments_fk_in_datasets_table.py
@@ -0,0 +1,67 @@
+"""add cascading delete on experiments fk in datasets table
+
+Revision ID: 93bd306463bd
+Revises: acf3f17fdcc7
+Create Date: 2024-04-12 10:13:01.319032
+
+"""
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = "93bd306463bd"
+down_revision = "acf3f17fdcc7"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    dialect_name = op.get_context().dialect.name
+    if dialect_name == "sqlite":
+        # Only way to drop unnamed fk constraint in sqllite
+        # See https://alembic.sqlalchemy.org/en/latest/batch.html#dropping-unnamed-or-named-foreign-key-constraints
+        with op.batch_alter_table(
+            "datasets",
+            schema=None,
+            naming_convention={
+                "fk": "fk_%(table_name)s_%(column_0_name)s_%(referred_table_name)s",
+            },
+        ) as batch_op:
+            batch_op.drop_constraint(
+                "fk_datasets_experiment_id_experiments", type_="foreignkey"
+            )
+            # Need to explicitly name the fk constraint with batch alter table
+            batch_op.create_foreign_key(
+                "fk_datasets_experiment_id_experiments",
+                "experiments",
+                ["experiment_id"],
+                ["experiment_id"],
+                ondelete="CASCADE",
+            )
+    else:
+        if dialect_name == "postgresql":
+            fk_constraint_name = "datasets_experiment_id_fkey"
+        elif dialect_name == "mysql":
+            fk_constraint_name = "datasets_ibfk_1"
+        elif dialect_name == "mssql":
+            # mssql fk constraint name at acf3f17fdcc7
+            fk_constraint_name = "FK__datasets__experi__6477ECF3"
+
+        # don't use batch alter table here so `create_foreign_key()` can be
+        # called with `None` as the `constraint_name` and have it set
+        # automatically based on the conventions of the sql flavor
+        op.drop_constraint(fk_constraint_name, "datasets", type_="foreignkey")
+        op.create_foreign_key(
+            None,
+            "datasets",
+            "experiments",
+            ["experiment_id"],
+            ["experiment_id"],
+            ondelete="CASCADE",
+        )
+
+
+def downgrade():
+    pass
diff --git a/mlflow/store/tracking/dbmodels/models.py b/mlflow/store/tracking/dbmodels/models.py
--- a/mlflow/store/tracking/dbmodels/models.py
+++ b/mlflow/store/tracking/dbmodels/models.py
@@ -491,7 +491,7 @@ class SqlDataset(Base):
     Dataset UUID: `String` (limit 36 characters). Defined as *Non-null* in schema.
     Part of *Primary Key* for ``datasets`` table.
     """
-    experiment_id = Column(Integer, ForeignKey("experiments.experiment_id"))
+    experiment_id = Column(Integer, ForeignKey("experiments.experiment_id", ondelete="CASCADE"))
     """
     Experiment ID to which this dataset belongs: *Foreign Key* into ``experiments`` table.
     """
@@ -555,7 +555,11 @@ class SqlInput(Base):
     __tablename__ = "inputs"
     __table_args__ = (
         PrimaryKeyConstraint(
-            "source_type", "source_id", "destination_type", "destination_id", name="inputs_pk"
+            "source_type",
+            "source_id",
+            "destination_type",
+            "destination_id",
+            name="inputs_pk",
         ),
         Index(f"index_{__tablename__}_input_uuid", "input_uuid"),
         Index(


Test patch:
diff --git a/tests/db/schemas/mssql.sql b/tests/db/schemas/mssql.sql
--- a/tests/db/schemas/mssql.sql
+++ b/tests/db/schemas/mssql.sql
@@ -53,7 +53,7 @@ CREATE TABLE datasets (
 	dataset_schema VARCHAR COLLATE "SQL_Latin1_General_CP1_CI_AS",
 	dataset_profile VARCHAR COLLATE "SQL_Latin1_General_CP1_CI_AS",
 	CONSTRAINT dataset_pk PRIMARY KEY (experiment_id, name, digest),
-	CONSTRAINT "FK__datasets__experi__6477ECF3" FOREIGN KEY(experiment_id) REFERENCES experiments (experiment_id)
+	CONSTRAINT "FK__datasets__experi__693CA210" FOREIGN KEY(experiment_id) REFERENCES experiments (experiment_id) ON DELETE CASCADE
 )
 
 
diff --git a/tests/db/schemas/mysql.sql b/tests/db/schemas/mysql.sql
--- a/tests/db/schemas/mysql.sql
+++ b/tests/db/schemas/mysql.sql
@@ -54,7 +54,7 @@ CREATE TABLE datasets (
 	dataset_schema TEXT,
 	dataset_profile MEDIUMTEXT,
 	PRIMARY KEY (experiment_id, name, digest),
-	CONSTRAINT datasets_ibfk_1 FOREIGN KEY(experiment_id) REFERENCES experiments (experiment_id)
+	CONSTRAINT datasets_ibfk_1 FOREIGN KEY(experiment_id) REFERENCES experiments (experiment_id) ON DELETE CASCADE
 )
 
 
diff --git a/tests/db/schemas/postgresql.sql b/tests/db/schemas/postgresql.sql
--- a/tests/db/schemas/postgresql.sql
+++ b/tests/db/schemas/postgresql.sql
@@ -55,7 +55,7 @@ CREATE TABLE datasets (
 	dataset_schema TEXT,
 	dataset_profile TEXT,
 	CONSTRAINT dataset_pk PRIMARY KEY (experiment_id, name, digest),
-	CONSTRAINT datasets_experiment_id_fkey FOREIGN KEY(experiment_id) REFERENCES experiments (experiment_id)
+	CONSTRAINT datasets_experiment_id_fkey FOREIGN KEY(experiment_id) REFERENCES experiments (experiment_id) ON DELETE CASCADE
 )
 
 
diff --git a/tests/db/schemas/sqlite.sql b/tests/db/schemas/sqlite.sql
--- a/tests/db/schemas/sqlite.sql
+++ b/tests/db/schemas/sqlite.sql
@@ -56,7 +56,7 @@ CREATE TABLE datasets (
 	dataset_schema TEXT,
 	dataset_profile TEXT,
 	CONSTRAINT dataset_pk PRIMARY KEY (experiment_id, name, digest),
-	FOREIGN KEY(experiment_id) REFERENCES experiments (experiment_id)
+	CONSTRAINT fk_datasets_experiment_id_experiments FOREIGN KEY(experiment_id) REFERENCES experiments (experiment_id) ON DELETE CASCADE
 )
 
 
diff --git a/tests/resources/db/latest_schema.sql b/tests/resources/db/latest_schema.sql
--- a/tests/resources/db/latest_schema.sql
+++ b/tests/resources/db/latest_schema.sql
@@ -56,7 +56,7 @@ CREATE TABLE datasets (
 	dataset_schema TEXT,
 	dataset_profile TEXT,
 	CONSTRAINT dataset_pk PRIMARY KEY (experiment_id, name, digest),
-	FOREIGN KEY(experiment_id) REFERENCES experiments (experiment_id)
+	CONSTRAINT fk_datasets_experiment_id_experiments FOREIGN KEY(experiment_id) REFERENCES experiments (experiment_id) ON DELETE CASCADE
 )
 
 
diff --git a/tests/test_cli.py b/tests/test_cli.py
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -18,7 +18,9 @@
 import mlflow
 from mlflow import pyfunc
 from mlflow.cli import doctor, gc, server
+from mlflow.data import numpy_dataset
 from mlflow.entities import ViewType
+from mlflow.entities.dataset_input import DatasetInput
 from mlflow.exceptions import MlflowException
 from mlflow.server import handlers
 from mlflow.store.tracking.file_store import FileStore
@@ -27,7 +29,11 @@
 from mlflow.utils.rest_utils import augmented_raise_for_status
 from mlflow.utils.time import get_current_time_millis
 
-from tests.helper_functions import PROTOBUF_REQUIREMENT, get_safe_port, pyfunc_serve_and_score_model
+from tests.helper_functions import (
+    PROTOBUF_REQUIREMENT,
+    get_safe_port,
+    pyfunc_serve_and_score_model,
+)
 from tests.tracking.integration_test_utils import _await_server_up_or_die
 
 
@@ -265,7 +271,14 @@ def test_mlflow_gc_file_store_passing_explicit_run_ids(file_store):
     run = _create_run_in_store(store)
     store.delete_run(run.info.run_uuid)
     subprocess.check_output(
-        ["mlflow", "gc", "--backend-store-uri", file_store[1], "--run-ids", run.info.run_uuid]
+        [
+            "mlflow",
+            "gc",
+            "--backend-store-uri",
+            file_store[1],
+            "--run-ids",
+            run.info.run_uuid,
+        ]
     )
     runs = store.search_runs(experiment_ids=["0"], filter_string="", run_view_type=ViewType.ALL)
     assert len(runs) == 0
@@ -278,7 +291,14 @@ def test_mlflow_gc_not_deleted_run(file_store):
     run = _create_run_in_store(store)
     with pytest.raises(subprocess.CalledProcessError, match=r".+"):
         subprocess.check_output(
-            ["mlflow", "gc", "--backend-store-uri", file_store[1], "--run-ids", run.info.run_uuid]
+            [
+                "mlflow",
+                "gc",
+                "--backend-store-uri",
+                file_store[1],
+                "--run-ids",
+                run.info.run_uuid,
+            ]
         )
     runs = store.search_runs(experiment_ids=["0"], filter_string="", run_view_type=ViewType.ALL)
     assert len(runs) == 1
@@ -376,13 +396,30 @@ def invoke_gc(*args):
     store.delete_experiment(exp_id_5)
     with pytest.raises(MlflowException, match=r"Experiments .+ can be deleted."):
         invoke_gc(
-            "--backend-store-uri", uri, "--experiment-ids", exp_id_5, "--older-than", "10d10h10m10s"
+            "--backend-store-uri",
+            uri,
+            "--experiment-ids",
+            exp_id_5,
+            "--older-than",
+            "10d10h10m10s",
         )
     experiments = store.search_experiments(view_type=ViewType.ALL)
     assert sorted([e.experiment_id for e in experiments]) == sorted(
         [exp_id_5, store.DEFAULT_EXPERIMENT_ID]
     )
 
+    exp_id_6 = store.create_experiment("6")
+    run_id_2 = store.create_run(exp_id_6, user_id="user", start_time=1, tags=[], run_name="2")
+    run_id_2_datasets = [
+        DatasetInput(dataset=numpy_dataset.from_numpy(np.array([1, 2, 3]))._to_mlflow_entity())
+    ]
+    store.log_inputs(run_id_2.info.run_id, datasets=run_id_2_datasets)
+    store.delete_experiment(exp_id_6)
+    invoke_gc("--backend-store-uri", uri, "--experiment-ids", exp_id_6)
+    assert sorted([e.experiment_id for e in experiments]) == sorted(
+        [exp_id_5, store.DEFAULT_EXPERIMENT_ID]
+    )
+
 
 @pytest.mark.parametrize(
     "enable_mlserver",