Repository: optuna/optuna

Issue:
Fix mypy errors when using SQLAlchemy v2.0.0.
### Motivation

Due to the 2nd major release of SQLAlchemy, mypy checks are failed. We added the version constraint at https://github.com/optuna/optuna/pull/4364 to avoid other PRs being blocked from merging.

### Suggestion

Fix mypy errors when using `SQLAlchemy==2.0.0` and remove a version constraint added at https://github.com/optuna/optuna/pull/4364/files.

### Additional context (optional)

See https://docs.sqlalchemy.org/en/14/changelog/migration_20.html for details.


Discussion:
Could this be the reason for mypy to get stuck or am I experiencing a different issue?

My mypy gets stuck on `optuna.samplers._tpe._erf`. 
Running `mypy -p optuna -v` prints a lot of logs, to finally get stuck here:
```
LOG:  Processing 4 queued fresh SCCs
LOG:  Processing SCC singleton (optuna.samplers._tpe._erf) as inherently stale --> does not finish
```

Note, I am running sqlalchemy 2.0.0 and optuna 3.1.0.
@nielsuit227 I couldn't reproduce it, so it is probably a different issue. My packages' versions are as follows:

```
- sqlalchemy 2.0.0
- mypy: 0.991
- mypy-extensions: 0.4.3
- optuna: 3.1.0 and latest 
```
Running the same versions, on Python 3.10.8. Problem occurs on the following minimal reproducable sample:
```python
import optuna

z = optuna.samplers.TPESampler()
```
Sqlalchemy does provide some [plugin](https://docs.sqlalchemy.org/en/20/orm/extensions/mypy.html) for better type checking. I suggest using (and modifying code with) this.
@gen740 I afraid that the document you linked states the plugin is deprecated in SQLAlchemy 2.0.
> The SQLAlchemy Mypy plugin, while it has technically never left the “alpha” stage, should now be considered as deprecated in SQLAlchemy 2.0, even though it is still necessary for full Mypy support when using SQLAlchemy 1.4.

Resolving patch:
diff --git a/optuna/storages/_rdb/models.py b/optuna/storages/_rdb/models.py
--- a/optuna/storages/_rdb/models.py
+++ b/optuna/storages/_rdb/models.py
@@ -8,7 +8,6 @@
 from sqlalchemy import asc
 from sqlalchemy import case
 from sqlalchemy import CheckConstraint
-from sqlalchemy import Column
 from sqlalchemy import DateTime
 from sqlalchemy import desc
 from sqlalchemy import Enum
@@ -32,6 +31,14 @@
     # TODO(c-bata): Remove this after dropping support for SQLAlchemy v1.3 or prior.
     from sqlalchemy.ext.declarative import declarative_base
 
+try:
+    from sqlalchemy.orm import mapped_column
+
+    _Column = mapped_column
+except ImportError:
+    # TODO(Shinichi): Remove this after dropping support for SQLAlchemy<2.0.
+    from sqlalchemy import Column as _Column  # type: ignore[assignment]
+
 
 # Don't modify this version number anymore.
 # The schema management functionality has been moved to alembic.
@@ -49,8 +56,10 @@
 
 class StudyModel(BaseModel):
     __tablename__ = "studies"
-    study_id = Column(Integer, primary_key=True)
-    study_name = Column(String(MAX_INDEXED_STRING_LENGTH), index=True, unique=True, nullable=False)
+    study_id = _Column(Integer, primary_key=True)
+    study_name = _Column(
+        String(MAX_INDEXED_STRING_LENGTH), index=True, unique=True, nullable=False
+    )
 
     @classmethod
     def find_or_raise_by_id(
@@ -85,10 +94,10 @@ def find_or_raise_by_name(cls, study_name: str, session: orm.Session) -> "StudyM
 class StudyDirectionModel(BaseModel):
     __tablename__ = "study_directions"
     __table_args__: Any = (UniqueConstraint("study_id", "objective"),)
-    study_direction_id = Column(Integer, primary_key=True)
-    direction = Column(Enum(StudyDirection), nullable=False)
-    study_id = Column(Integer, ForeignKey("studies.study_id"), nullable=False)
-    objective = Column(Integer, nullable=False)
+    study_direction_id = _Column(Integer, primary_key=True)
+    direction = _Column(Enum(StudyDirection), nullable=False)
+    study_id = _Column(Integer, ForeignKey("studies.study_id"), nullable=False)
+    objective = _Column(Integer, nullable=False)
 
     study = orm.relationship(
         StudyModel, backref=orm.backref("directions", cascade="all, delete-orphan")
@@ -102,10 +111,10 @@ def where_study_id(cls, study_id: int, session: orm.Session) -> List["StudyDirec
 class StudyUserAttributeModel(BaseModel):
     __tablename__ = "study_user_attributes"
     __table_args__: Any = (UniqueConstraint("study_id", "key"),)
-    study_user_attribute_id = Column(Integer, primary_key=True)
-    study_id = Column(Integer, ForeignKey("studies.study_id"))
-    key = Column(String(MAX_INDEXED_STRING_LENGTH))
-    value_json = Column(Text())
+    study_user_attribute_id = _Column(Integer, primary_key=True)
+    study_id = _Column(Integer, ForeignKey("studies.study_id"))
+    key = _Column(String(MAX_INDEXED_STRING_LENGTH))
+    value_json = _Column(Text())
 
     study = orm.relationship(
         StudyModel, backref=orm.backref("user_attributes", cascade="all, delete-orphan")
@@ -134,10 +143,10 @@ def where_study_id(
 class StudySystemAttributeModel(BaseModel):
     __tablename__ = "study_system_attributes"
     __table_args__: Any = (UniqueConstraint("study_id", "key"),)
-    study_system_attribute_id = Column(Integer, primary_key=True)
-    study_id = Column(Integer, ForeignKey("studies.study_id"))
-    key = Column(String(MAX_INDEXED_STRING_LENGTH))
-    value_json = Column(Text())
+    study_system_attribute_id = _Column(Integer, primary_key=True)
+    study_id = _Column(Integer, ForeignKey("studies.study_id"))
+    key = _Column(String(MAX_INDEXED_STRING_LENGTH))
+    value_json = _Column(Text())
 
     study = orm.relationship(
         StudyModel, backref=orm.backref("system_attributes", cascade="all, delete-orphan")
@@ -165,15 +174,15 @@ def where_study_id(
 
 class TrialModel(BaseModel):
     __tablename__ = "trials"
-    trial_id = Column(Integer, primary_key=True)
+    trial_id = _Column(Integer, primary_key=True)
     # No `UniqueConstraint` is put on the `number` columns although it in practice is constrained
     # to be unique. This is to reduce code complexity as table-level locking would be required
     # otherwise. See https://github.com/optuna/optuna/pull/939#discussion_r387447632.
-    number = Column(Integer)
-    study_id = Column(Integer, ForeignKey("studies.study_id"), index=True)
-    state = Column(Enum(TrialState), nullable=False)
-    datetime_start = Column(DateTime)
-    datetime_complete = Column(DateTime)
+    number = _Column(Integer)
+    study_id = _Column(Integer, ForeignKey("studies.study_id"), index=True)
+    state = _Column(Enum(TrialState), nullable=False)
+    datetime_start = _Column(DateTime)
+    datetime_complete = _Column(DateTime)
 
     study = orm.relationship(
         StudyModel, backref=orm.backref("trials", cascade="all, delete-orphan")
@@ -273,10 +282,10 @@ def count_past_trials(self, session: orm.Session) -> int:
 class TrialUserAttributeModel(BaseModel):
     __tablename__ = "trial_user_attributes"
     __table_args__: Any = (UniqueConstraint("trial_id", "key"),)
-    trial_user_attribute_id = Column(Integer, primary_key=True)
-    trial_id = Column(Integer, ForeignKey("trials.trial_id"))
-    key = Column(String(MAX_INDEXED_STRING_LENGTH))
-    value_json = Column(Text())
+    trial_user_attribute_id = _Column(Integer, primary_key=True)
+    trial_id = _Column(Integer, ForeignKey("trials.trial_id"))
+    key = _Column(String(MAX_INDEXED_STRING_LENGTH))
+    value_json = _Column(Text())
 
     trial = orm.relationship(
         TrialModel, backref=orm.backref("user_attributes", cascade="all, delete-orphan")
@@ -305,10 +314,10 @@ def where_trial_id(
 class TrialSystemAttributeModel(BaseModel):
     __tablename__ = "trial_system_attributes"
     __table_args__: Any = (UniqueConstraint("trial_id", "key"),)
-    trial_system_attribute_id = Column(Integer, primary_key=True)
-    trial_id = Column(Integer, ForeignKey("trials.trial_id"))
-    key = Column(String(MAX_INDEXED_STRING_LENGTH))
-    value_json = Column(Text())
+    trial_system_attribute_id = _Column(Integer, primary_key=True)
+    trial_id = _Column(Integer, ForeignKey("trials.trial_id"))
+    key = _Column(String(MAX_INDEXED_STRING_LENGTH))
+    value_json = _Column(Text())
 
     trial = orm.relationship(
         TrialModel, backref=orm.backref("system_attributes", cascade="all, delete-orphan")
@@ -337,11 +346,11 @@ def where_trial_id(
 class TrialParamModel(BaseModel):
     __tablename__ = "trial_params"
     __table_args__: Any = (UniqueConstraint("trial_id", "param_name"),)
-    param_id = Column(Integer, primary_key=True)
-    trial_id = Column(Integer, ForeignKey("trials.trial_id"))
-    param_name = Column(String(MAX_INDEXED_STRING_LENGTH))
-    param_value = Column(Float(precision=FLOAT_PRECISION))
-    distribution_json = Column(Text())
+    param_id = _Column(Integer, primary_key=True)
+    trial_id = _Column(Integer, ForeignKey("trials.trial_id"))
+    param_name = _Column(String(MAX_INDEXED_STRING_LENGTH))
+    param_value = _Column(Float(precision=FLOAT_PRECISION))
+    distribution_json = _Column(Text())
 
     trial = orm.relationship(
         TrialModel, backref=orm.backref("params", cascade="all, delete-orphan")
@@ -408,11 +417,11 @@ class TrialValueType(enum.Enum):
 
     __tablename__ = "trial_values"
     __table_args__: Any = (UniqueConstraint("trial_id", "objective"),)
-    trial_value_id = Column(Integer, primary_key=True)
-    trial_id = Column(Integer, ForeignKey("trials.trial_id"), nullable=False)
-    objective = Column(Integer, nullable=False)
-    value = Column(Float(precision=FLOAT_PRECISION), nullable=True)
-    value_type = Column(Enum(TrialValueType), nullable=False)
+    trial_value_id = _Column(Integer, primary_key=True)
+    trial_id = _Column(Integer, ForeignKey("trials.trial_id"), nullable=False)
+    objective = _Column(Integer, nullable=False)
+    value = _Column(Float(precision=FLOAT_PRECISION), nullable=True)
+    value_type = _Column(Enum(TrialValueType), nullable=False)
 
     trial = orm.relationship(
         TrialModel, backref=orm.backref("values", cascade="all, delete-orphan")
@@ -474,11 +483,11 @@ class TrialIntermediateValueType(enum.Enum):
 
     __tablename__ = "trial_intermediate_values"
     __table_args__: Any = (UniqueConstraint("trial_id", "step"),)
-    trial_intermediate_value_id = Column(Integer, primary_key=True)
-    trial_id = Column(Integer, ForeignKey("trials.trial_id"), nullable=False)
-    step = Column(Integer, nullable=False)
-    intermediate_value = Column(Float(precision=FLOAT_PRECISION), nullable=True)
-    intermediate_value_type = Column(Enum(TrialIntermediateValueType), nullable=False)
+    trial_intermediate_value_id = _Column(Integer, primary_key=True)
+    trial_id = _Column(Integer, ForeignKey("trials.trial_id"), nullable=False)
+    step = _Column(Integer, nullable=False)
+    intermediate_value = _Column(Float(precision=FLOAT_PRECISION), nullable=True)
+    intermediate_value_type = _Column(Enum(TrialIntermediateValueType), nullable=False)
 
     trial = orm.relationship(
         TrialModel, backref=orm.backref("intermediate_values", cascade="all, delete-orphan")
@@ -541,9 +550,9 @@ def where_trial_id(
 class TrialHeartbeatModel(BaseModel):
     __tablename__ = "trial_heartbeats"
     __table_args__: Any = (UniqueConstraint("trial_id"),)
-    trial_heartbeat_id = Column(Integer, primary_key=True)
-    trial_id = Column(Integer, ForeignKey("trials.trial_id"), nullable=False)
-    heartbeat = Column(DateTime, nullable=False, default=func.current_timestamp())
+    trial_heartbeat_id = _Column(Integer, primary_key=True)
+    trial_id = _Column(Integer, ForeignKey("trials.trial_id"), nullable=False)
+    heartbeat = _Column(DateTime, nullable=False, default=func.current_timestamp())
 
     trial = orm.relationship(
         TrialModel, backref=orm.backref("heartbeats", cascade="all, delete-orphan")
@@ -560,12 +569,11 @@ class VersionInfoModel(BaseModel):
     __tablename__ = "version_info"
     # setting check constraint to ensure the number of rows is at most 1
     __table_args__: Any = (CheckConstraint("version_info_id=1"),)
-    version_info_id = Column(Integer, primary_key=True, autoincrement=False, default=1)
-    schema_version = Column(Integer)
-    library_version = Column(String(MAX_VERSION_LENGTH))
+    version_info_id = _Column(Integer, primary_key=True, autoincrement=False, default=1)
+    schema_version = _Column(Integer)
+    library_version = _Column(String(MAX_VERSION_LENGTH))
 
     @classmethod
-    def find(cls, session: orm.Session) -> "VersionInfoModel":
+    def find(cls, session: orm.Session) -> Optional["VersionInfoModel"]:
         version_info = session.query(cls).one_or_none()
-
         return version_info
diff --git a/optuna/storages/_rdb/storage.py b/optuna/storages/_rdb/storage.py
--- a/optuna/storages/_rdb/storage.py
+++ b/optuna/storages/_rdb/storage.py
@@ -12,6 +12,7 @@
 from typing import Container
 from typing import Dict
 from typing import Generator
+from typing import Iterable
 from typing import List
 from typing import Optional
 from typing import Sequence
@@ -802,6 +803,9 @@ def _get_trials(
             )
 
             if states is not None:
+                # This assertion is for type checkers, since `states` is required to be Container
+                # in the base class while `models.TrialModel.state.in_` requires Iterable.
+                assert isinstance(states, Iterable)
                 query = query.filter(models.TrialModel.state.in_(states))
 
             trial_ids = query.all()
@@ -989,6 +993,7 @@ def _get_stale_trial_ids(self, study_id: int) -> List[int]:
 
         with _create_scoped_session(self.scoped_session, True) as session:
             current_heartbeat = session.execute(sqlalchemy.func.now()).scalar()
+            assert current_heartbeat is not None
             # Added the following line to prevent mixing of timezone-aware and timezone-naive
             # `datetime` in PostgreSQL. See
             # https://github.com/optuna/optuna/pull/2190#issuecomment-766605088 for details


Test patch:
diff --git a/tests/storages_tests/rdb_tests/test_models.py b/tests/storages_tests/rdb_tests/test_models.py
--- a/tests/storages_tests/rdb_tests/test_models.py
+++ b/tests/storages_tests/rdb_tests/test_models.py
@@ -130,6 +130,7 @@ def test_default_datetime(session: Session) -> None:
 
         trial_model = session.query(TrialModel).first()
 
+        assert trial_model is not None
         assert trial_model.datetime_start is None
         assert trial_model.datetime_complete is None
 
diff --git a/tests/storages_tests/rdb_tests/test_storage.py b/tests/storages_tests/rdb_tests/test_storage.py
--- a/tests/storages_tests/rdb_tests/test_storage.py
+++ b/tests/storages_tests/rdb_tests/test_storage.py
@@ -34,6 +34,7 @@ def test_init() -> None:
     session = storage.scoped_session()
 
     version_info = session.query(VersionInfoModel).first()
+    assert version_info is not None
     assert version_info.schema_version == SCHEMA_VERSION
     assert version_info.library_version == optuna.version.__version__
 
@@ -55,6 +56,7 @@ def test_init() -> None:
 def test_init_url_template() -> None:
     with NamedTemporaryFilePool(suffix="{SCHEMA_VERSION}") as tf:
         storage = RDBStorage("sqlite:///" + tf.name)
+        assert storage.engine.url.database is not None
         assert storage.engine.url.database.endswith(str(SCHEMA_VERSION))