code: def upgrade() -> None:
      bind = op.get_bind()
      inspector = sa.inspect(bind)
      tables = inspector.get_table_names()
  
      assert "trial_params" in tables
  
      session = orm.Session(bind=bind)
      try:
          distributions: List[BaseDistribution] = []
          for distribution in session.query(TrialParamModel).yield_per(
              BATCH_SIZE
          ):
              distribution.distribution_json = migrate_new_distribution(
                  distribution.distribution_json,
              )
              distributions.append(distribution)
  
              if len(distributions) == BATCH_SIZE:
                  persist(session, distributions)
                  distributions = []
  
          persist(session, distributions)
  
      except SQLAlchemyError as e:
          session.rollback()
          raise e
      finally:
          session.close()
  
imports: from typing import Any
  from typing import List
  import sqlalchemy
  from alembic import op
  from sqlalchemy import Column
  from sqlalchemy import DateTime
  from sqlalchemy import Enum
  from sqlalchemy import Float
  from sqlalchemy import ForeignKey
  from sqlalchemy import Integer
  from sqlalchemy import orm
  from sqlalchemy import String
  from sqlalchemy import Text
  from sqlalchemy import UniqueConstraint
  from sqlalchemy.exc import SQLAlchemyError
  from optuna.distributions import _convert_old_distribution_to_new_distribution
  from optuna.distributions import BaseDistribution
  from optuna.distributions import DiscreteUniformDistribution
  from optuna.distributions import distribution_to_json
  from optuna.distributions import FloatDistribution
  from optuna.distributions import IntDistribution
  from optuna.distributions import IntLogUniformDistribution
  from optuna.distributions import IntUniformDistribution
  from optuna.distributions import json_to_distribution
  from optuna.distributions import LogUniformDistribution
  from optuna.distributions import UniformDistribution
  from optuna.trial import TrialState
label_desc: Use image node_id 4 for calling a global function with example usage: upgrade() without return types