Repository: mlflow/mlflow

Issue:
[BUG] MLflow REST API incorrectly reports "database down" as BAD_REQUEST (HTTP 400) error
### 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 can contribute a fix for this bug independently.

### MLflow version

- Client: 1.30.0
- Tracking server: 1.30.0


### System information

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


### Describe the problem

I'm testing the behavior MLflow REST server when the database server (e.g. Postgres) goes down. I discovered that the server returns an HTTP 400 **client** error, which is not retried by the MLflow client. Hence, a job using MLflow fails immediately rather than retrying.

The incorrect code is [here](https://github.com/mlflow/mlflow/blob/v1.30.0/mlflow/store/db/utils.py#L105).

I propose adding another `except` block above this one, resulting in:
```
        except sqlalchemy.exc.OperationalError as e:
            session.rollback()
            raise MlflowException(message=e, error_code= INTERNAL_ERROR)
        except sqlalchemy.exc.SQLAlchemyError as e:
            session.rollback()
            raise MlflowException(message=e, error_code=BAD_REQUEST)
```

[PEP 249 – Python Database API Specification v2.0 | peps.python.org](https://peps.python.org/pep-0249/#operationalerror) [Core Exceptions — SQLAlchemy 1.4 Documentation] describes `OperationalError` as:
> Exception raised for errors that are related to the database’s operation and not necessarily under the control of the programmer, e.g. an unexpected disconnect occurs, the data source name is not found, a transaction could not be processed, a memory allocation error occurred during processing, etc.

This exception type is primarily for errors that occur at runtime, relating to the operational environment of the application, rather than a bug in the application (or something the MLflow user / client did wrong).



### Tracking information

_No response_

### Code to reproduce issue

```
import mlflow

experiment_name = f"foobar"
mlflow.set_experiment(experiment_name)
with mlflow.start_run():
    pass
```

### Stack trace

```
Traceback (most recent call last):
  File "/usr/local/lib/python3.7/site-packages/mlflow/server/handlers.py", line 455, in wrapper
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.7/site-packages/mlflow/server/handlers.py", line 528, in wrapper
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.7/site-packages/mlflow/server/handlers.py", line 711, in _update_run
    run_id, request_message.status, request_message.end_time, request_message.run_name
  File "/usr/local/lib/python3.7/site-packages/mlflow/store/tracking/sqlalchemy_store.py", line 650, in update_run_info
    return run.info
  File "/usr/local/lib/python3.7/contextlib.py", line 130, in __exit__
    self.gen.throw(type, value, traceback)
  File "/usr/local/lib/python3.7/site-packages/mlflow/store/db/utils.py", line 105, in make_managed_session
    raise MlflowException(message=e, error_code=BAD_REQUEST)
mlflow.exceptions.MlflowException: (psycopg2.OperationalError) could not connect to server: Connection refused
        Is the server running on host "127.0.0.1" and accepting
        TCP/IP connections on port 5432?
```

### 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/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
- [ ] `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:
Sounds good to me! Please let us know when you file the PR :) 
Done! Created PR #7240. Are there pre-existing tests for this area of the code?

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
@@ -23,7 +23,7 @@
 
 from mlflow.exceptions import MlflowException
 from mlflow.store.tracking.dbmodels.initial_models import Base as InitialBase
-from mlflow.protos.databricks_pb2 import BAD_REQUEST, INTERNAL_ERROR
+from mlflow.protos.databricks_pb2 import BAD_REQUEST, INTERNAL_ERROR, TEMPORARILY_UNAVAILABLE
 from mlflow.store.db.db_types import SQLITE
 from mlflow.environment_variables import (
     MLFLOW_SQLALCHEMYSTORE_POOL_SIZE,
@@ -100,6 +100,13 @@ def make_managed_session():
         except MlflowException:
             session.rollback()
             raise
+        except sqlalchemy.exc.OperationalError as e:
+            session.rollback()
+            _logger.exception(
+                "SQLAlchemy database error. The following exception is caught.\n%s",
+                e,
+            )
+            raise MlflowException(message=e, error_code=TEMPORARILY_UNAVAILABLE)
         except sqlalchemy.exc.SQLAlchemyError as e:
             session.rollback()
             raise MlflowException(message=e, error_code=BAD_REQUEST)


Test patch:
diff --git a/tests/db/test_tracking_operations.py b/tests/db/test_tracking_operations.py
--- a/tests/db/test_tracking_operations.py
+++ b/tests/db/test_tracking_operations.py
@@ -1,7 +1,14 @@
+import os
+import sqlite3
+import uuid
+from unittest import mock
+
 import pytest
+import sqlalchemy.dialects.sqlite.pysqlite
 
 import mlflow
 from mlflow import MlflowClient
+from mlflow.tracking._tracking_service.utils import _TRACKING_URI_ENV_VAR
 
 
 pytestmark = pytest.mark.notrackingurimock
@@ -41,3 +48,114 @@ def test_set_run_status_to_killed():
         pass
     client = MlflowClient()
     client.set_terminated(run_id=run.info.run_id, status="KILLED")
+
+
+@mock.patch("mlflow.store.db.utils._logger.exception")
+def test_database_operational_error(exception, monkeypatch):
+    # This test is specifically designed to force errors with SQLite. Skip it if
+    # using a non-SQLite backend.
+    if not os.environ[_TRACKING_URI_ENV_VAR].startswith("sqlite"):
+        pytest.skip("Only works on SQLite")
+
+    # This test patches parts of SQLAlchemy and sqlite3.dbapi to simulate a
+    # SQLAlchemy OperationalError. PEP 249 describes OperationalError as:
+    #
+    # > Exception raised for errors that are related to the database’s operation
+    # > and not necessarily under the control of the programmer, e.g. an
+    # > unexpected disconnect occurs, the data source name is not found, a
+    # > transaction could not be processed, a memory allocation error occurred
+    # > during processing, etc.
+    #
+    # These errors are typically transient and can be resolved by retrying the
+    # operation, hence MLflow has different handling for them as compared to
+    # the more generic exception type, SQLAlchemyError.
+    #
+    # This is particularly important for REST clients, where
+    # TEMPORARILY_UNAVAILABLE triggers MLflow REST clients to retry the request,
+    # whereas BAD_REQUEST does not.
+    api_module = None
+    old_connect = None
+
+    # Depending on the version of SQLAlchemy, the function we need to patch is
+    # either called "dbapi" (sqlalchemy<2.0) or "import_dbapi"
+    # (sqlalchemy>=2.0).
+    for dialect_attr in ["dbapi", "import_dbapi"]:
+        if hasattr(sqlalchemy.dialects.sqlite.pysqlite.SQLiteDialect_pysqlite, dialect_attr):
+            break
+    else:
+        raise AssertionError("Could not find dbapi attribute on SQLiteDialect_pysqlite")
+
+    old_dbapi = getattr(sqlalchemy.dialects.sqlite.pysqlite.SQLiteDialect_pysqlite, dialect_attr)
+
+    class ConnectionWrapper:
+        """Wraps a sqlite3.Connection object."""
+
+        def __init__(self, conn):
+            self.conn = conn
+
+        def __getattr__(self, name):
+            return getattr(self.conn, name)
+
+        def cursor(self):
+            """Return a wrapped SQLite cursor."""
+            return CursorWrapper(self.conn.cursor())
+
+    class CursorWrapper:
+        """Wraps a sqlite3.Cursor object."""
+
+        def __init__(self, cursor):
+            self.cursor = cursor
+
+        def __getattr__(self, name):
+            return getattr(self.cursor, name)
+
+        def execute(self, *args, **kwargs):
+            """Wraps execute(), simulating sporadic OperationalErrors."""
+            if (
+                len(args) >= 2
+                and "test_database_operational_error_1667938883_param" in args[1]
+                and "test_database_operational_error_1667938883_value" in args[1]
+            ):
+                # Simulate a database error
+                raise sqlite3.OperationalError("test")
+            return self.cursor.execute(*args, **kwargs)
+
+    def connect(*args, **kwargs):
+        """Wraps sqlite3.dbapi.connect(), returning a wrapped connection."""
+        global connect_counter
+        conn = old_connect(*args, **kwargs)  # pylint: disable=not-callable
+        return ConnectionWrapper(conn)
+
+    def dbapi(*args, **kwargs):
+        """Wraps SQLiteDialect_pysqlite.dbapi(), returning patched dbapi."""
+        nonlocal api_module, old_connect
+        if api_module is None:
+            # Only patch the first time dbapi() is called, to avoid recursion.
+            api_module = old_dbapi(*args, **kwargs)
+            old_connect = api_module.connect
+            monkeypatch.setattr(api_module, "connect", connect)
+        return api_module
+
+    monkeypatch.setattr(
+        sqlalchemy.dialects.sqlite.pysqlite.SQLiteDialect_pysqlite, dialect_attr, dbapi
+    )
+
+    # Create and use a unique tracking URI for this test. This avoids an issue
+    # where an earlier test has already created and cached a SQLAlchemy engine
+    # (i.e. database connections), preventing our error-throwing monkeypatches
+    # from being called.
+    monkeypatch.setitem(
+        os.environ, _TRACKING_URI_ENV_VAR, f"{os.environ[_TRACKING_URI_ENV_VAR]}-{uuid.uuid4().hex}"
+    )
+    with pytest.raises(mlflow.MlflowException, match=r"sqlite3\.OperationalError"):
+        with mlflow.start_run():
+            # This statement will fail with an OperationalError.
+            mlflow.log_param(
+                "test_database_operational_error_1667938883_param",
+                "test_database_operational_error_1667938883_value",
+            )
+    # Verify that the error handling was executed.
+    assert any(
+        "SQLAlchemy database error" in str(call) and "sqlite3.OperationalError" in str(call)
+        for call in exception.mock_calls
+    )
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
@@ -31,6 +31,7 @@
     BAD_REQUEST,
     RESOURCE_DOES_NOT_EXIST,
     INVALID_PARAMETER_VALUE,
+    TEMPORARILY_UNAVAILABLE,
 )
 from mlflow.store.tracking import SEARCH_MAX_RESULTS_DEFAULT
 from mlflow.store.db.utils import (
@@ -1143,15 +1144,27 @@ def test_log_null_param(self):
         tval = None
         param = entities.Param(tkey, tval)
 
+        dialect = self.store._get_dialect()
         regex = {
             SQLITE: r"NOT NULL constraint failed",
             POSTGRES: r"null value in column .+ of relation .+ violates not-null constrain",
             MYSQL: r"Column .+ cannot be null",
             MSSQL: r"Cannot insert the value NULL into column .+, table .+",
-        }[self.store._get_dialect()]
+        }[dialect]
         with pytest.raises(MlflowException, match=regex) as exception_context:
             self.store.log_param(run.info.run_id, param)
-        assert exception_context.value.error_code == ErrorCode.Name(BAD_REQUEST)
+        if dialect != MYSQL:
+            assert exception_context.value.error_code == ErrorCode.Name(BAD_REQUEST)
+        else:
+            # Some MySQL client packages (and there are several available, e.g.
+            # PyMySQL, mysqlclient, mysql-connector-python... reports some
+            # errors, including NULL constraint violations, as a SQLAlchemy
+            # OperationalError, even though they should be reported as a more
+            # generic SQLAlchemyError. If that is fixed, we can remove this
+            # special case.
+            assert exception_context.value.error_code == ErrorCode.Name(
+                BAD_REQUEST
+            ) or exception_context.value.error_code == ErrorCode.Name(TEMPORARILY_UNAVAILABLE)
 
     def test_log_param_max_length_value(self):
         run = self._run_factory()