Avoid to use features that will be removed in SQLAlchemy v2.0

.github/workflows/tests-storage.yml

@@ -23,6 +23,7 @@ jobs:
     strategy:
       matrix:
         python-version: ['3.7', '3.8', '3.9', '3.10', '3.11']
+        sqlalchemy-version: ['1.4.46', '2.0.0rc2']
 
     services:
       mysql:
@@ -93,6 +94,9 @@ jobs:
 
         pip install --progress-bar off .[test]
         pip install --progress-bar off .[optional]
+        pip install --progress-bar off sqlalchemy==$SQLALCHEMY_VERSION
+      env:
+        SQLALCHEMY_VERSION: ${{ matrix.sqlalchemy-version }}
 
     - name: Install DB bindings
       run: |
@@ -102,6 +106,7 @@ jobs:
       run: |
         pytest tests/storages_tests/test_with_server.py
       env:
+        SQLALCHEMY_WARN_20: 1
         OMP_NUM_THREADS: 1
         TEST_DB_URL: mysql+pymysql://user:test@127.0.0.1/optunatest
 

.github/workflows/tests.yml

@@ -21,6 +21,7 @@ jobs:
     strategy:
       matrix:
         python-version: ['3.7', '3.8', '3.9', '3.10', '3.11']
+        sqlalchemy-version: ['1.4.46', '2.0.0rc2']
 
     services:
       redis:
@@ -65,13 +66,20 @@ jobs:
 
         pip install --progress-bar off .[test]
         pip install --progress-bar off .[optional]
+        pip install --progress-bar off sqlalchemy==$SQLALCHEMY_VERSION
+      env:
+        SQLALCHEMY_VERSION: ${{ matrix.sqlalchemy-version }}
 
     - name: Scheduled tests
       if:  ${{ github.event_name == 'schedule' || github.event_name == 'workflow_dispatch' }}
       run: |
         pytest tests -m "not integration"
+      env:
+        SQLALCHEMY_WARN_20: 1
 
     - name: Tests
       if:  ${{ github.event_name != 'schedule' }}
       run: |
         pytest tests -m "not integration and not slow"
+      env:
+        SQLALCHEMY_WARN_20: 1

optuna/storages/_rdb/alembic/versions/v1.3.0.a.py

@@ -11,9 +11,14 @@
 import sqlalchemy as sa
 
 from sqlalchemy.exc import SQLAlchemyError
-from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy import orm
 
+try:
+    from sqlalchemy.orm import declarative_base
+except ImportError:
+    # TODO(c-bata): Remove this after dropping support for SQLAlchemy v1.3 or prior.
+    from sqlalchemy.ext.declarative import declarative_base
+
 # revision identifiers, used by Alembic.
 revision = "v1.3.0.a"
 down_revision = "v1.2.0.a"

optuna/storages/_rdb/alembic/versions/v2.4.0.a.py

@@ -16,11 +16,16 @@
 from sqlalchemy import Integer
 from sqlalchemy import UniqueConstraint
 from sqlalchemy.exc import SQLAlchemyError
-from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy import orm
 
 from optuna.study import StudyDirection
 
+try:
+    from sqlalchemy.orm import declarative_base
+except ImportError:
+    # TODO(c-bata): Remove this after dropping support for SQLAlchemy v1.3 or prior.
+    from sqlalchemy.ext.declarative import declarative_base
+
 
 # revision identifiers, used by Alembic.
 revision = "v2.4.0.a"

optuna/storages/_rdb/alembic/versions/v3.0.0.a.py

@@ -21,7 +21,6 @@
 from sqlalchemy import Text
 from sqlalchemy import UniqueConstraint
 from sqlalchemy.exc import SQLAlchemyError
-from sqlalchemy.ext.declarative import declarative_base
 
 from optuna.distributions import _convert_old_distribution_to_new_distribution
 from optuna.distributions import BaseDistribution
@@ -36,6 +35,12 @@
 from optuna.distributions import UniformDistribution
 from optuna.trial import TrialState
 
+try:
+    from sqlalchemy.orm import declarative_base
+except ImportError:
+    # TODO(c-bata): Remove this after dropping support for SQLAlchemy v1.3 or prior.
+    from sqlalchemy.ext.declarative import declarative_base
+
 
 # revision identifiers, used by Alembic.
 revision = "v3.0.0.a"

optuna/storages/_rdb/alembic/versions/v3.0.0.b.py

@@ -14,9 +14,14 @@
 from sqlalchemy import Float
 from sqlalchemy import ForeignKey
 from sqlalchemy import Integer
-from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.orm import Session
 
+try:
+    from sqlalchemy.orm import declarative_base
+except ImportError:
+    # TODO(c-bata): Remove this after dropping support for SQLAlchemy v1.3 or prior.
+    from sqlalchemy.ext.declarative import declarative_base
+
 
 # revision identifiers, used by Alembic.
 revision = "v3.0.0.b"

optuna/storages/_rdb/alembic/versions/v3.0.0.c.py

@@ -11,11 +11,16 @@
 from alembic import op
 import sqlalchemy as sa
 from sqlalchemy.exc import SQLAlchemyError
-from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy import orm
 from typing import Optional
 from typing import Tuple
 
+try:
+    from sqlalchemy.orm import declarative_base
+except ImportError:
+    # TODO(c-bata): Remove this after dropping support for SQLAlchemy v1.3 or prior.
+    from sqlalchemy.ext.declarative import declarative_base
+
 
 # revision identifiers, used by Alembic.
 revision = "v3.0.0.c"

optuna/storages/_rdb/alembic/versions/v3.0.0.d.py

@@ -11,11 +11,16 @@
 from alembic import op
 import sqlalchemy as sa
 from sqlalchemy.exc import SQLAlchemyError
-from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy import orm
 from typing import Optional
 from typing import Tuple
 
+try:
+    from sqlalchemy.orm import declarative_base
+except ImportError:
+    # TODO(c-bata): Remove this after dropping support for SQLAlchemy v1.3 or prior.
+    from sqlalchemy.ext.declarative import declarative_base
+
 
 # revision identifiers, used by Alembic.
 revision = "v3.0.0.d"

optuna/storages/_rdb/models.py

@@ -19,13 +19,19 @@
 from sqlalchemy import String
 from sqlalchemy import Text
 from sqlalchemy import UniqueConstraint
-from sqlalchemy.ext.declarative import declarative_base
 
 from optuna import distributions
 from optuna.study._study_direction import StudyDirection
 from optuna.trial import TrialState
 
 
+try:
+    from sqlalchemy.orm import declarative_base
+except ImportError:
+    # TODO(c-bata): Remove this after dropping support for SQLAlchemy v1.3 or prior.
+    from sqlalchemy.ext.declarative import declarative_base
+
+
 # Don't modify this version number anymore.
 # The schema management functionality has been moved to alembic.
 SCHEMA_VERSION = 12

optuna/storages/_rdb/storage.py

@@ -1098,9 +1098,11 @@ def _init_alembic(self) -> None:
 
     def _set_alembic_revision(self, revision: str) -> None:
 
-        context = alembic_migration.MigrationContext.configure(self.engine.connect())
-        script = self._create_alembic_script()
-        context.stamp(script, revision)
+        connection = self.engine.connect()
+        context = alembic_migration.MigrationContext.configure(connection)
+        with connection.begin():
+            script = self._create_alembic_script()
+            context.stamp(script, revision)
 
     def check_table_schema_compatibility(self) -> None: