Repository: mlflow/mlflow

Issue:
[BUG] `_all_tables_exist` is out of date 
### 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)

### Willingness to contribute

No. I cannot contribute a bug fix at this time.

### MLflow version

- Client: 2.4.1
- Tracking server: 2.4.1
These I am using, but seems the related code is the same in latest version.

### System information

- **Python version**: 3.10

### Describe the problem

We observed more DB connections after server start than in previous versions (2.1, 2.3).

I think it might be caused by check _all_tables_exist from mlflow/store/db/utils.py never passing and therefore db upgrade always runs. There are some new tables in the DB: datasets, input_tags, inputs, which are not considered in the check.


### Other info / logs

~$mlflow db upgrade postgresql://mlflow_user:mlflow@localhost/mlflow_db                                                                                                        
2023/07/11 09:39:01 INFO mlflow.store.db.utils: Updating database tables
INFO  [alembic.runtime.migration] Context impl PostgresqlImpl.
INFO  [alembic.runtime.migration] Will assume transactional DDL.
~$mlflow server --backend-store-uri postgresql://mlflow_user:mlflow@localhost/mlflow_db --default-artifact-root sftp://mlflow_user@127.0.0.0:~/mlflow/mlruns -h 0.0.0.0 -p 8000
mlflow/store/db/utils.py: _all_tables_exist:
DB:       ['datasets', 'experiment_tags', 'experiments', 'input_tags', 'inputs', 'latest_metrics', 'metrics', 'model_version_tags', 'model_versions', 'params', 'registered_model_aliases', 'registered_model_tags', 'registered_models', 'runs', 'tags']
expected: ['experiment_tags', 'experiments', 'latest_metrics', 'metrics', 'model_version_tags', 'model_versions', 'params', 'registered_model_aliases', 'registered_model_tags', 'registered_models', 'runs', 'tags']
2023/07/11 09:39:06 INFO mlflow.store.db.utils: Creating initial MLflow database tables...
2023/07/11 09:39:06 INFO mlflow.store.db.utils: Updating database tables
INFO  [alembic.runtime.migration] Context impl PostgresqlImpl.
INFO  [alembic.runtime.migration] Will assume transactional DDL.
INFO  [alembic.runtime.migration] Context impl PostgresqlImpl.
INFO  [alembic.runtime.migration] Will assume transactional DDL.
mlflow/store/db/utils.py: _all_tables_exist:
DB:       ['datasets', 'experiment_tags', 'experiments', 'input_tags', 'inputs', 'latest_metrics', 'metrics', 'model_version_tags', 'model_versions', 'params', 'registered_model_aliases', 'registered_model_tags', 'registered_models', 'runs', 'tags']
expected: ['experiment_tags', 'experiments', 'latest_metrics', 'metrics', 'model_version_tags', 'model_versions', 'params', 'registered_model_aliases', 'registered_model_tags', 'registered_models', 'runs', 'tags']
2023/07/11 09:39:06 INFO mlflow.store.db.utils: Creating initial MLflow database tables...
2023/07/11 09:39:06 INFO mlflow.store.db.utils: Updating database tables
INFO  [alembic.runtime.migration] Context impl PostgresqlImpl.
INFO  [alembic.runtime.migration] Will assume transactional DDL.
[2023-07-11 09:39:06 +0200] [42489] [INFO] Starting gunicorn 20.1.0
[2023-07-11 09:39:06 +0200] [42489] [INFO] Listening at: http://0.0.0.0:8000 (42489)
[2023-07-11 09:39:06 +0200] [42489] [INFO] Using worker: sync
[2023-07-11 09:39:06 +0200] [42490] [INFO] Booting worker with pid: 42490
[2023-07-11 09:39:06 +0200] [42491] [INFO] Booting worker with pid: 42491
[2023-07-11 09:39:06 +0200] [42492] [INFO] Booting worker with pid: 42492
[2023-07-11 09:39:06 +0200] [42493] [INFO] Booting worker with pid: 42493




Discussion:

Resolving patch:
diff --git a/mlflow/store/db/utils.py b/mlflow/store/db/utils.py
--- a/mlflow/store/db/utils.py
+++ b/mlflow/store/db/utils.py
@@ -31,6 +31,9 @@
     SqlTag,
     SqlExperimentTag,
     SqlLatestMetric,
+    SqlDataset,
+    SqlInput,
+    SqlInputTag,
 )
 from mlflow.store.model_registry.dbmodels.models import (
     SqlRegisteredModel,
@@ -79,6 +82,9 @@ def _all_tables_exist(engine):
         SqlRegisteredModelTag.__tablename__,
         SqlModelVersionTag.__tablename__,
         SqlRegisteredModelAlias.__tablename__,
+        SqlDataset.__tablename__,
+        SqlInput.__tablename__,
+        SqlInputTag.__tablename__,
     }
 
 


Test patch: