Repository: mlflow/mlflow

Issue:
[BUG] INTERNAL_SERVER_ERROR when using MSSQL as backend store
### Willingness to contribute
The MLflow Community encourages bug fix contributions. Would you or another member of your organization be willing to contribute a fix for this bug to the MLflow code base?

- [ ] Yes. I can contribute a fix for this bug independently.
- [x] Yes. I would be willing to contribute a fix for this bug with guidance from the MLflow community.
- [ ] No. I cannot contribute a bug fix at this time.

### System information
- **Have I written custom code (as opposed to using a stock example script provided in MLflow)**: no
- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**: Mac OS Big Sur
- **MLflow installed from (source or binary)**: binary
- **MLflow version (run ``mlflow --version``)**: version 1.23.1
- **Python version**: Python 3.9.9
- **npm version, if running the dev UI**: -

### Describe the problem
When running a tracking ui with `mlflow ui --backend-store-uri  <connection_str>` and specifying the backend store a MSSQL database, there is a problem when trying to access the UI. 
<img width="1274" alt="image" src="https://user-images.githubusercontent.com/63718157/153061210-2ca9fdc5-1090-4468-a99d-02dfa7f6f227.png">

### Code to reproduce issue
Start the tracking UI with a backend store a MSSQL database.
Or do:
````
mlflow.set_tracking_uri(connection_str_to_mssql_db)
mlflow.search_runs(experiment_ids=[0], order_by= ['attributes.start_time DESC'])
````


### Other info / logs
In the logs, there is the following error:
````
2022/02/08 14:30:14 ERROR mlflow.server: Exception on /ajax-api/2.0/preview/mlflow/runs/search [POST]
Traceback (most recent call last):
  File "/Users/dianacarvalho/opt/anaconda3/envs/vpdc-complete/lib/python3.9/site-packages/sqlalchemy/engine/base.py", line 1802, in _execute_context
    self.dialect.do_execute(
  File "/Users/dianacarvalho/opt/anaconda3/envs/vpdc-complete/lib/python3.9/site-packages/sqlalchemy/engine/default.py", line 732, in do_execute
    cursor.execute(statement, parameters)
pyodbc.ProgrammingError: ('42000', '[42000] [Microsoft][ODBC Driver 17 for SQL Server][SQL Server]ORDER BY items must appear in the select list if SELECT DISTINCT is specified. (145) (SQLExecDirectW)')

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/Users/dianacarvalho/opt/anaconda3/envs/vpdc-complete/lib/python3.9/site-packages/mlflow/store/db/utils.py", line 79, in make_managed_session
    yield session
  File "/Users/dianacarvalho/opt/anaconda3/envs/vpdc-complete/lib/python3.9/site-packages/mlflow/store/tracking/sqlalchemy_store.py", line 788, in _search_runs
    query.distinct()
  File "/Users/dianacarvalho/opt/anaconda3/envs/vpdc-complete/lib/python3.9/site-packages/sqlalchemy/orm/query.py", line 2759, in all
    return self._iter().all()
  File "/Users/dianacarvalho/opt/anaconda3/envs/vpdc-complete/lib/python3.9/site-packages/sqlalchemy/orm/query.py", line 2894, in _iter
    result = self.session.execute(
  File "/Users/dianacarvalho/opt/anaconda3/envs/vpdc-complete/lib/python3.9/site-packages/sqlalchemy/orm/session.py", line 1692, in execute
    result = conn._execute_20(statement, params or {}, execution_options)
  File "/Users/dianacarvalho/opt/anaconda3/envs/vpdc-complete/lib/python3.9/site-packages/sqlalchemy/engine/base.py", line 1614, in _execute_20
    return meth(self, args_10style, kwargs_10style, execution_options)
  File "/Users/dianacarvalho/opt/anaconda3/envs/vpdc-complete/lib/python3.9/site-packages/sqlalchemy/sql/elements.py", line 325, in _execute_on_connection
    return connection._execute_clauseelement(
  File "/Users/dianacarvalho/opt/anaconda3/envs/vpdc-complete/lib/python3.9/site-packages/sqlalchemy/engine/base.py", line 1481, in _execute_clauseelement
    ret = self._execute_context(
  File "/Users/dianacarvalho/opt/anaconda3/envs/vpdc-complete/lib/python3.9/site-packages/sqlalchemy/engine/base.py", line 1845, in _execute_context
    self._handle_dbapi_exception(
  File "/Users/dianacarvalho/opt/anaconda3/envs/vpdc-complete/lib/python3.9/site-packages/sqlalchemy/engine/base.py", line 2026, in _handle_dbapi_exception
    util.raise_(
  File "/Users/dianacarvalho/opt/anaconda3/envs/vpdc-complete/lib/python3.9/site-packages/sqlalchemy/util/compat.py", line 207, in raise_
    raise exception
  File "/Users/dianacarvalho/opt/anaconda3/envs/vpdc-complete/lib/python3.9/site-packages/sqlalchemy/engine/base.py", line 1802, in _execute_context
    self.dialect.do_execute(
  File "/Users/dianacarvalho/opt/anaconda3/envs/vpdc-complete/lib/python3.9/site-packages/sqlalchemy/engine/default.py", line 732, in do_execute
    cursor.execute(statement, parameters)
sqlalchemy.exc.ProgrammingError: (pyodbc.ProgrammingError) ('42000', '[42000] [Microsoft][ODBC Driver 17 for SQL Server][SQL Server]ORDER BY items must appear in the select list if SELECT DISTINCT is specified. (145) (SQLExecDirectW)')
[SQL: SELECT DISTINCT runs.run_uuid AS runs_run_uuid, runs.name AS runs_name, runs.source_type AS runs_source_type, runs.source_name AS runs_source_name, runs.entry_point_name AS runs_entry_point_name, runs.user_id AS runs_user_id, runs.status AS runs_status, runs.start_time AS runs_start_time, runs.end_time AS runs_end_time, runs.source_version AS runs_source_version, runs.lifecycle_stage AS runs_lifecycle_stage, runs.artifact_uri AS runs_artifact_uri, runs.experiment_id AS runs_experiment_id, CASE WHEN (runs.start_time IS NULL) THEN ? ELSE ? END AS anon_1 
FROM runs 
WHERE runs.experiment_id IN (?) AND runs.lifecycle_stage IN (?) ORDER BY CASE WHEN (runs.start_time IS NULL) THEN ? ELSE ? END, runs.start_time DESC, runs.run_uuid
 OFFSET ? ROWS
 FETCH FIRST ? ROWS ONLY]
[parameters: (1, 0, '0', 'active', 1, 0, 0, 100)]
(Background on this error at: https://sqlalche.me/e/14/f405)

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/Users/dianacarvalho/opt/anaconda3/envs/vpdc-complete/lib/python3.9/site-packages/mlflow/server/handlers.py", line 238, in wrapper
    return func(*args, **kwargs)
  File "/Users/dianacarvalho/opt/anaconda3/envs/vpdc-complete/lib/python3.9/site-packages/mlflow/server/handlers.py", line 298, in wrapper
    return func(*args, **kwargs)
  File "/Users/dianacarvalho/opt/anaconda3/envs/vpdc-complete/lib/python3.9/site-packages/mlflow/server/handlers.py", line 549, in _search_runs
    run_entities = _get_tracking_store().search_runs(
  File "/Users/dianacarvalho/opt/anaconda3/envs/vpdc-complete/lib/python3.9/site-packages/mlflow/store/tracking/abstract_store.py", line 242, in search_runs
    runs, token = self._search_runs(
  File "/Users/dianacarvalho/opt/anaconda3/envs/vpdc-complete/lib/python3.9/site-packages/mlflow/store/tracking/sqlalchemy_store.py", line 802, in _search_runs
    next_page_token = compute_next_token(len(runs))
  File "/Users/dianacarvalho/opt/anaconda3/envs/vpdc-complete/lib/python3.9/contextlib.py", line 137, in __exit__
    self.gen.throw(typ, value, traceback)
  File "/Users/dianacarvalho/opt/anaconda3/envs/vpdc-complete/lib/python3.9/site-packages/mlflow/store/db/utils.py", line 86, in make_managed_session
    raise MlflowException(message=e, error_code=INTERNAL_ERROR)
mlflow.exceptions.MlflowException: (pyodbc.ProgrammingError) ('42000', '[42000] [Microsoft][ODBC Driver 17 for SQL Server][SQL Server]ORDER BY items must appear in the select list if SELECT DISTINCT is specified. (145) (SQLExecDirectW)')
[SQL: SELECT DISTINCT runs.run_uuid AS runs_run_uuid, runs.name AS runs_name, runs.source_type AS runs_source_type, runs.source_name AS runs_source_name, runs.entry_point_name AS runs_entry_point_name, runs.user_id AS runs_user_id, runs.status AS runs_status, runs.start_time AS runs_start_time, runs.end_time AS runs_end_time, runs.source_version AS runs_source_version, runs.lifecycle_stage AS runs_lifecycle_stage, runs.artifact_uri AS runs_artifact_uri, runs.experiment_id AS runs_experiment_id, CASE WHEN (runs.start_time IS NULL) THEN ? ELSE ? END AS anon_1 
FROM runs 
WHERE runs.experiment_id IN (?) AND runs.lifecycle_stage IN (?) ORDER BY CASE WHEN (runs.start_time IS NULL) THEN ? ELSE ? END, runs.start_time DESC, runs.run_uuid
 OFFSET ? ROWS
 FETCH FIRST ? ROWS ONLY]
[parameters: (1, 0, '0', 'active', 1, 0, 0, 100)]
(Background on this error at: https://sqlalche.me/e/14/f405)

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/Users/dianacarvalho/opt/anaconda3/envs/vpdc-complete/lib/python3.9/site-packages/flask/app.py", line 2073, in wsgi_app
    response = self.full_dispatch_request()
  File "/Users/dianacarvalho/opt/anaconda3/envs/vpdc-complete/lib/python3.9/site-packages/flask/app.py", line 1518, in full_dispatch_request
    rv = self.handle_user_exception(e)
  File "/Users/dianacarvalho/opt/anaconda3/envs/vpdc-complete/lib/python3.9/site-packages/flask/app.py", line 1516, in full_dispatch_request
    rv = self.dispatch_request()
  File "/Users/dianacarvalho/opt/anaconda3/envs/vpdc-complete/lib/python3.9/site-packages/flask/app.py", line 1502, in dispatch_request
    return self.ensure_sync(self.view_functions[rule.endpoint])(**req.view_args)
  File "/Users/dianacarvalho/opt/anaconda3/envs/vpdc-complete/lib/python3.9/site-packages/mlflow/server/handlers.py", line 241, in wrapper
    response.set_data(e.serialize_as_json())
  File "/Users/dianacarvalho/opt/anaconda3/envs/vpdc-complete/lib/python3.9/site-packages/mlflow/exceptions.py", line 60, in serialize_as_json
    return json.dumps(exception_dict)
  File "/Users/dianacarvalho/opt/anaconda3/envs/vpdc-complete/lib/python3.9/json/__init__.py", line 231, in dumps
    return _default_encoder.encode(obj)
  File "/Users/dianacarvalho/opt/anaconda3/envs/vpdc-complete/lib/python3.9/json/encoder.py", line 199, in encode
    chunks = self.iterencode(o, _one_shot=True)
  File "/Users/dianacarvalho/opt/anaconda3/envs/vpdc-complete/lib/python3.9/json/encoder.py", line 257, in iterencode
    return _iterencode(o, 0)
  File "/Users/dianacarvalho/opt/anaconda3/envs/vpdc-complete/lib/python3.9/json/encoder.py", line 179, in default
    raise TypeError(f'Object of type {o.__class__.__name__} '
TypeError: Object of type ProgrammingError is not JSON serializable
````
### What component(s), interfaces, languages, and integrations does this bug affect?
Components 
- [ ] `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/projects`: MLproject format, project running backends
- [ ] `area/scoring`: MLflow Model server, model deployment tools, Spark UDFs
- [x] `area/server-infra`: MLflow Tracking server backend
- [ ] `area/tracking`: Tracking Service, tracking client APIs, autologging

Interface 
- [ ] `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




Discussion:
I believe this is a sqlalchemy bug instead so opened an issue [there](https://github.com/sqlalchemy/sqlalchemy/issues/7696) as well.

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
@@ -6,6 +6,7 @@
 import math
 import sqlalchemy
 import sqlalchemy.sql.expression as sql
+from sqlalchemy.future import select
 
 from mlflow.entities.lifecycle_stage import LifecycleStage
 from mlflow.store.tracking import SEARCH_MAX_RESULTS_THRESHOLD
@@ -475,13 +476,13 @@ def _get_eager_run_query_options():
                  run attributes when fetching a run: ``latest_metrics``, ``params``, and ``tags``.
         """
         return [
-            # Use a subquery load rather than a joined load in order to minimize the memory overhead
-            # of the eager loading procedure. For more information about relationship loading
-            # techniques, see https://docs.sqlalchemy.org/en/13/orm/
+            # Use a select in load rather than a joined load in order to minimize the memory
+            # overhead of the eager loading procedure. For more information about relationship
+            # loading techniques, see https://docs.sqlalchemy.org/en/13/orm/
             # loading_relationships.html#relationship-loading-techniques
-            sqlalchemy.orm.subqueryload(SqlRun.latest_metrics),
-            sqlalchemy.orm.subqueryload(SqlRun.params),
-            sqlalchemy.orm.subqueryload(SqlRun.tags),
+            sqlalchemy.orm.selectinload(SqlRun.latest_metrics),
+            sqlalchemy.orm.selectinload(SqlRun.params),
+            sqlalchemy.orm.selectinload(SqlRun.tags),
         ]
 
     def _check_run_is_active(self, run):
@@ -772,20 +773,20 @@ def compute_next_token(current_size):
             # ``run.to_mlflow_entity()``, so eager loading helps avoid additional database queries
             # that are otherwise executed at attribute access time under a lazy loading model.
             parsed_filters = SearchUtils.parse_search_filter(filter_string)
-            parsed_orderby, sorting_joins = _get_orderby_clauses(order_by, session)
+            cases_orderby, parsed_orderby, sorting_joins = _get_orderby_clauses(order_by, session)
 
-            query = session.query(SqlRun)
+            stmt = select(SqlRun, *cases_orderby)
             for j in _get_sqlalchemy_filter_clauses(parsed_filters, session):
-                query = query.join(j)
+                stmt = stmt.join(j)
             # using an outer join is necessary here because we want to be able to sort
             # on a column (tag, metric or param) without removing the lines that
             # do not have a value for this column (which is what inner join would do)
             for j in sorting_joins:
-                query = query.outerjoin(j)
+                stmt = stmt.outerjoin(j)
 
             offset = SearchUtils.parse_start_offset_from_page_token(page_token)
-            queried_runs = (
-                query.distinct()
+            stmt = (
+                stmt.distinct()
                 .options(*self._get_eager_run_query_options())
                 .filter(
                     SqlRun.experiment_id.in_(experiment_ids),
@@ -795,8 +796,8 @@ def compute_next_token(current_size):
                 .order_by(*parsed_orderby)
                 .offset(offset)
                 .limit(max_results)
-                .all()
             )
+            queried_runs = session.execute(stmt).scalars(SqlRun).all()
 
             runs = [run.to_mlflow_entity() for run in queried_runs]
             next_page_token = compute_next_token(len(runs))
@@ -920,6 +921,7 @@ def _get_orderby_clauses(order_by_list, session):
     ordering_joins = []
     clause_id = 0
     observed_order_by_clauses = set()
+    select_clauses = []
     # contrary to filters, it is not easily feasible to separately handle sorting
     # on attributes and on joined tables as we must keep all clauses in the same order
     if order_by_list:
@@ -956,24 +958,26 @@ def _get_orderby_clauses(order_by_list, session):
             # same main query, the CASE WHEN columns need to have unique names to
             # avoid ambiguity
             if SearchUtils.is_metric(key_type, "="):
-                clauses.append(
-                    sql.case(
-                        [
-                            # Ideally the use of "IS" is preferred here but owing to sqlalchemy
-                            # translation in MSSQL we are forced to use "=" instead.
-                            # These 2 options are functionally identical / unchanged because
-                            # the column (is_nan) is not nullable. However it could become an issue
-                            # if this precondition changes in the future.
-                            (subquery.c.is_nan == sqlalchemy.true(), 1),
-                            (order_value.is_(None), 1),
-                        ],
-                        else_=0,
-                    ).label("clause_%s" % clause_id)
-                )
+                case = sql.case(
+                    [
+                        # Ideally the use of "IS" is preferred here but owing to sqlalchemy
+                        # translation in MSSQL we are forced to use "=" instead.
+                        # These 2 options are functionally identical / unchanged because
+                        # the column (is_nan) is not nullable. However it could become an issue
+                        # if this precondition changes in the future.
+                        (subquery.c.is_nan == sqlalchemy.true(), 1),
+                        (order_value.is_(None), 1),
+                    ],
+                    else_=0,
+                ).label("clause_%s" % clause_id)
+
             else:  # other entities do not have an 'is_nan' field
-                clauses.append(
-                    sql.case([(order_value.is_(None), 1)], else_=0).label("clause_%s" % clause_id)
+                case = sql.case([(order_value.is_(None), 1)], else_=0).label(
+                    "clause_%s" % clause_id
                 )
+            clauses.append(case.name)
+            select_clauses.append(case)
+            select_clauses.append(order_value)
 
             if (key_type, key) in observed_order_by_clauses:
                 raise MlflowException(
@@ -989,4 +993,4 @@ def _get_orderby_clauses(order_by_list, session):
     if (SearchUtils._ATTRIBUTE_IDENTIFIER, SqlRun.start_time.key) not in observed_order_by_clauses:
         clauses.append(SqlRun.start_time.desc())
     clauses.append(SqlRun.run_uuid)
-    return clauses, ordering_joins
+    return select_clauses, clauses, ordering_joins


Test patch:
diff --git a/tests/db/run_checks.py b/tests/db/run_checks.py
--- a/tests/db/run_checks.py
+++ b/tests/db/run_checks.py
@@ -33,6 +33,9 @@ def run_logging_operations():
             python_model=MockModel(),
             registered_model_name="mock",
         )
+    runs = mlflow.search_runs(experiment_ids=["0"], order_by=["param.start_time DESC"])
+
+    run = mlflow.get_run(runs["run_id"][0])
 
     # Ensure the following migration scripts work correctly:
     # - cfd24bdc0731_update_run_status_constraint_with_killed.py
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
@@ -2005,11 +2005,11 @@ def test_get_orderby_clauses():
     store = SqlAlchemyStore("sqlite:///:memory:", ARTIFACT_URI)
     with store.ManagedSessionMaker() as session:
         # test that ['runs.start_time DESC', 'SqlRun.run_uuid'] is returned by default
-        parsed = [str(x) for x in _get_orderby_clauses([], session)[0]]
+        parsed = [str(x) for x in _get_orderby_clauses([], session)[1]]
         assert parsed == ["runs.start_time DESC", "SqlRun.run_uuid"]
 
         # test that the given 'start_time' replaces the default one ('runs.start_time DESC')
-        parsed = [str(x) for x in _get_orderby_clauses(["attribute.start_time ASC"], session)[0]]
+        parsed = [str(x) for x in _get_orderby_clauses(["attribute.start_time ASC"], session)[1]]
         assert "SqlRun.start_time" in parsed
         assert "SqlRun.start_time DESC" not in parsed
 
@@ -2030,7 +2030,11 @@ def test_get_orderby_clauses():
         # test that an exception is NOT raised when key types are different
         _get_orderby_clauses(["param.a", "metric.a", "tag.a"], session)
 
+        select_clause, parsed, _ = _get_orderby_clauses(["metric.a"], session)
+        select_clause = [str(x) for x in select_clause]
+        parsed = [str(x) for x in parsed]
         # test that "=" is used rather than "is" when comparing to True
-        parsed = [str(x) for x in _get_orderby_clauses(["metric.a"], session)[0]]
-        assert "is_nan = true" in parsed[0]
-        assert "value IS NULL" in parsed[0]
+        assert "is_nan = true" in select_clause[0]
+        assert "value IS NULL" in select_clause[0]
+        # test that clause name is in parsed
+        assert "clause_1" in parsed[0]