Repository: mlflow/mlflow

Issue:
[BUG] Tags not properly updated when updating run name in `SqlAlchemyStore`
### 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

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

### MLflow version

d38a78f8

### System information

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


### Describe the problem

related to #6932 and #7013 

When using postgresql as tracking uri, and when updating run name via `store.update_run_info()`, the `mlflow.user` is modified but not the expected `mlflow.runName`.

The problem does not occur for `FileStore`



### Tracking information

MLflow version: 1.29.1.dev0

### Code to reproduce issue

```python

import mlflow
from mlflow.tracking import MlflowClient
from mlflow.entities.run_status import RunStatus

if __name__ == "__main__":

    tracking_uri = "file:///C:/Users/gabrielfu/PycharmProjects/mlflow/mlruns"
    # tracking_uri = "postgresql+psycopg2://user:123456@localhost/mlflow_db"
    mlflow.set_tracking_uri(tracking_uri)
    print("Current tracking uri: {}".format(mlflow.get_tracking_uri()))

    # start a new run
    with mlflow.start_run() as run:
        run_id = run.info.run_id
        pass

    run = mlflow.get_run(run_id)
    print("before:")
    print("run:", run)
    print("==")

    def update(run_id, run_name):
        run = mlflow.get_run(run_id)
        client = MlflowClient()
        client._tracking_client.store.update_run_info(
            run_id=run_id,
            run_status=RunStatus.from_string(run.info.status),
            end_time=run.info.end_time,
            run_name=run_name,
        )

    update(run_id, "new_name")
    run = mlflow.get_run(run_id)
    print("after:")
    print("run:", run)
```

### Stack trace

### Running with FileStore
Output is as expected:
```
Current tracking uri: file:///C:/Users/gabrielfu/PycharmProjects/mlflow/mlruns
before:
run: <Run: data=<RunData: metrics={}, params={}, tags={'mlflow.runName': 'tasteful-dog-864',
 'mlflow.source.git.commit': '8d5d0a08c23383ef8b86c8fa67bdb12e4467304d',
 'mlflow.source.name': 'C:/Users/gabrielfu/PycharmProjects/mlflow/code/get_run.py',
 'mlflow.source.type': 'LOCAL',
 'mlflow.user': 'gabrielfu'}>, info=<RunInfo: artifact_uri='./mlruns/0/e7ee01bc1bb0411faca6b12ba394b1d2/artifacts', end_time=1665475081092, experiment_id='0', lifecycle_stage='active', run_id='e7ee01bc1bb0411faca6b12ba394b1d2', run_name='tasteful-dog-864', run_uuid='e7ee01bc1bb0411faca6b12ba394b1d2', start_time=1665475081052, status='FINISHED', user_id='gabrielfu'>>
==
after:
run: <Run: data=<RunData: metrics={}, params={}, tags={'mlflow.runName': 'new_name',
 'mlflow.source.git.commit': '8d5d0a08c23383ef8b86c8fa67bdb12e4467304d',
 'mlflow.source.name': 'C:/Users/gabrielfu/PycharmProjects/mlflow/code/get_run.py',
 'mlflow.source.type': 'LOCAL',
 'mlflow.user': 'gabrielfu'}>, info=<RunInfo: artifact_uri='./mlruns/0/e7ee01bc1bb0411faca6b12ba394b1d2/artifacts', end_time=1665475081092, experiment_id='0', lifecycle_stage='active', run_id='e7ee01bc1bb0411faca6b12ba394b1d2', run_name='new_name', run_uuid='e7ee01bc1bb0411faca6b12ba394b1d2', start_time=1665475081052, status='FINISHED', user_id='gabrielfu'>>
```

### Running with SqlAlchemyStore
If I change to `tracking_uri = "postgresql+psycopg2://user:123456@localhost/mlflow_db"`, then `mlflow.user` tag got changed but not `mlflow.runName`.
```
Current tracking uri: postgresql+psycopg2://user:123456@localhost/mlflow_db
before:
run: <Run: data=<RunData: metrics={}, params={}, tags={'mlflow.runName': 'fun-goat-459',
 'mlflow.source.git.commit': '8d5d0a08c23383ef8b86c8fa67bdb12e4467304d',
 'mlflow.source.name': 'C:/Users/gabrielfu/PycharmProjects/mlflow/code/get_run.py',
 'mlflow.source.type': 'LOCAL',
 'mlflow.user': 'gabrielfu'}>, info=<RunInfo: artifact_uri='./mlruns/0/b09a343ca34844288775f65462cb0352/artifacts', end_time=1665475607550, experiment_id='0', lifecycle_stage='active', run_id='b09a343ca34844288775f65462cb0352', run_name='fun-goat-459', run_uuid='b09a343ca34844288775f65462cb0352', start_time=1665475607522, status='FINISHED', user_id='gabrielfu'>>
==
after:
run: <Run: data=<RunData: metrics={}, params={}, tags={'mlflow.runName': 'fun-goat-459',
 'mlflow.source.git.commit': '8d5d0a08c23383ef8b86c8fa67bdb12e4467304d',
 'mlflow.source.name': 'C:/Users/gabrielfu/PycharmProjects/mlflow/code/get_run.py',
 'mlflow.source.type': 'LOCAL',
 'mlflow.user': 'new_name'}>, info=<RunInfo: artifact_uri='./mlruns/0/b09a343ca34844288775f65462cb0352/artifacts', end_time=1665475607550, experiment_id='0', lifecycle_stage='active', run_id='b09a343ca34844288775f65462cb0352', run_name='new_name', run_uuid='b09a343ca34844288775f65462cb0352', start_time=1665475607522, status='FINISHED', user_id='gabrielfu'>>
```

### Other info / logs

_No response_

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

- [ ] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/build`: Build and test infrastructure for MLflow
- [ ] `area/docs`: MLflow documentation pages
- [ ] `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/pipelines`: Pipelines, Pipeline APIs, Pipeline configs, Pipeline Templates
- [ ] `area/projects`: MLproject format, project running backends
- [ ] `area/scoring`: MLflow Model server, model deployment tools, Spark UDFs
- [ ] `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
- [ ] `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:
@gabrielfu Thanks for reporting the issue. I was able to reproduce it.

I found `_try_get_run_tag` is causing this issue.

https://github.com/mlflow/mlflow/blob/b98bda05aa0257c327dd157bcc42251a406bbfef/mlflow/store/tracking/sqlalchemy_store.py#L638-L650

Line 643 needs to be fixed as follows:

```diff
-         .filter(SqlTag.run_uuid == run_id and SqlTag.key == tagKey)
+         .filter(SqlTag.run_uuid == run_id, SqlTag.key == tagKey)
```

### Why `SqlTag.run_uuid == run_id and SqlTag.key == tagKey` doesn't work:

`bool(SqlTag.run_uuid == run_id)` returns `False`, which means `SqlTag.run_uuid == run_id and SqlTag.key == tagKey` is equivalent to `SqlTag.run_uuid == run_id` since `and` returns the first falsey value. Since the tags table is only filtered by run ID, `tags` contains the all the tags associated with the specified run. `tags[0]` could be `mlflow.user`, not `mlflow.runName`.
@gabrielfu Please let me know if you would like to file a PR for this issue.
@harupy Thanks for the swift reply. I would like to submit a PR :) Now that you have found the cause, I should have enough information to work on it.

Are there any test cases related to this that need to change?
~I think just fixing line 643 is fine :)~
Can we add a test for `_try_get_run_tag`?
Can we also fix `_try_get_run_tag` like this using `one_or_none`?

```diff
diff --git a/mlflow/store/tracking/sqlalchemy_store.py b/mlflow/store/tracking/sqlalchemy_store.py
index 38cf6e14c..4fbf387a6 100644
--- a/mlflow/store/tracking/sqlalchemy_store.py
+++ b/mlflow/store/tracking/sqlalchemy_store.py
@@ -637,13 +637,12 @@ class SqlAlchemyStore(AbstractStore):
 
     def _try_get_run_tag(self, session, run_id, tagKey, eager=False):
         query_options = self._get_eager_run_query_options() if eager else []
-        tags = (
+        return (
             session.query(SqlTag)
             .options(*query_options)
             .filter(SqlTag.run_uuid == run_id and SqlTag.key == tagKey)
-            .all()
+            .one_or_none()
         )
-        return None if not tags else tags[0]
```
definitely

Resolving patch:
diff --git a/mlflow/store/tracking/sqlalchemy_store.py b/mlflow/store/tracking/sqlalchemy_store.py
--- a/mlflow/store/tracking/sqlalchemy_store.py
+++ b/mlflow/store/tracking/sqlalchemy_store.py
@@ -648,13 +648,12 @@ def update_run_info(self, run_id, run_status, end_time, run_name):
 
     def _try_get_run_tag(self, session, run_id, tagKey, eager=False):
         query_options = self._get_eager_run_query_options() if eager else []
-        tags = (
+        return (
             session.query(SqlTag)
             .options(*query_options)
-            .filter(SqlTag.run_uuid == run_id and SqlTag.key == tagKey)
-            .all()
+            .filter(SqlTag.run_uuid == run_id, SqlTag.key == tagKey)
+            .one_or_none()
         )
-        return None if not tags else tags[0]
 
     def get_run(self, run_id):
         with self.ManagedSessionMaker() as session:


Test patch:
diff --git a/tests/store/tracking/test_sqlalchemy_store.py b/tests/store/tracking/test_sqlalchemy_store.py
--- a/tests/store/tracking/test_sqlalchemy_store.py
+++ b/tests/store/tracking/test_sqlalchemy_store.py
@@ -2601,6 +2601,23 @@ def test_search_runs_keep_all_runs_when_sorting(self):
         )
         assert len(run_results) == 2
 
+    def test_try_get_run_tag(self):
+        run = self._run_factory()
+        self.store.set_tag(run.info.run_id, entities.RunTag("k1", "v1"))
+        self.store.set_tag(run.info.run_id, entities.RunTag("k2", "v2"))
+
+        with self.store.ManagedSessionMaker() as session:
+            tag = self.store._try_get_run_tag(session, run.info.run_id, "k0")
+            self.assertIsNone(tag)
+
+            tag = self.store._try_get_run_tag(session, run.info.run_id, "k1")
+            self.assertEqual(tag.key, "k1")
+            self.assertEqual(tag.value, "v1")
+
+            tag = self.store._try_get_run_tag(session, run.info.run_id, "k2")
+            self.assertEqual(tag.key, "k2")
+            self.assertEqual(tag.value, "v2")
+
 
 def test_sqlalchemy_store_behaves_as_expected_with_inmemory_sqlite_db():
     store = SqlAlchemyStore("sqlite:///:memory:", ARTIFACT_URI)