content: """Simplify model version links [ec6307720f92].
  
  Revision ID: ec6307720f92
  Revises: 0.70.0
  Create Date: 2024-11-06 16:16:43.344569
  
  """
  
  import sqlalchemy as sa
  from alembic import op
  
  # revision identifiers, used by Alembic.
  revision = "ec6307720f92"
  down_revision = "0.70.0"
  branch_labels = None
  depends_on = None
  
  
  def _migrate_artifact_type() -> None:
      """Migrate the artifact type."""
      meta = sa.MetaData()
      meta.reflect(
          bind=op.get_bind(),
          only=("artifact_version", "model_versions_artifacts"),
      )
      artifact_version_table = sa.Table("artifact_version", meta)
      model_versions_artifacts = sa.Table("model_versions_artifacts", meta)
  
      connection = op.get_bind()
  
      query = sa.select(
          model_versions_artifacts.c.artifact_version_id,
          model_versions_artifacts.c.is_model_artifact,
          model_versions_artifacts.c.is_deployment_artifact,
      )
      result = connection.execute(query)
  
      updated = set()
      updates = []
      for (
          artifact_version_id,
          is_model_artifact,
          is_deployment_artifact,
      ) in result:
          if artifact_version_id in updated:
              # If an artifact was a model artifact in one model version, and a
              # deployment artifact in another model version, we only update it
              # once.
              continue
  
          if is_model_artifact:
              updated.add(artifact_version_id)
              updates.append(
                  {"id_": artifact_version_id, "type": "ModelArtifact"}
              )
          elif is_deployment_artifact:
              updated.add(artifact_version_id)
              updates.append(
                  {"id_": artifact_version_id, "type": "ServiceArtifact"}
              )
  
      if updates:
          connection.execute(
              sa.update(artifact_version_table).where(
                  artifact_version_table.c.id == sa.bindparam("id_")
              ),
              updates,
          )
  
  
  def upgrade() -> None:
      """Upgrade database schema and/or data, creating a new revision."""
      # ### commands auto generated by Alembic - please adjust! ###
  
      _migrate_artifact_type()
  
      with op.batch_alter_table(
          "model_versions_artifacts", schema=None
      ) as batch_op:
          batch_op.drop_constraint(
              "fk_model_versions_artifacts_workspace_id_workspace",
              type_="foreignkey",
          )
          batch_op.drop_constraint(
              "fk_model_versions_artifacts_user_id_user", type_="foreignkey"
          )
          batch_op.drop_constraint(
              "fk_model_versions_artifacts_model_id_model", type_="foreignkey"
          )
          batch_op.drop_column("user_id")
          batch_op.drop_column("model_id")
          batch_op.drop_column("is_deployment_artifact")
          batch_op.drop_column("workspace_id")
          batch_op.drop_column("is_model_artifact")
  
      with op.batch_alter_table("model_versions_runs", schema=None) as batch_op:
          batch_op.drop_constraint(
              "fk_model_versions_runs_model_id_model", type_="foreignkey"
          )
          batch_op.drop_constraint(
              "fk_model_versions_runs_workspace_id_workspace", type_="foreignkey"
          )
          batch_op.drop_constraint(
              "fk_model_versions_runs_user_id_user", type_="foreignkey"
          )
          batch_op.drop_column("model_id")
          batch_op.drop_column("workspace_id")
          batch_op.drop_column("user_id")
  
      # ### end Alembic commands ###
  
  
  def downgrade() -> None:
      """Downgrade database schema and/or data back to the previous revision.
  
      Raises:
          NotImplementedError: Downgrade is not supported for this migration.
      """
      raise NotImplementedError("Downgrade is not supported for this migration.")
  