content: """Artifact Version Table [a91762e6be36].
  
  Revision ID: a91762e6be36
  Revises: 0.50.0
  Create Date: 2023-11-25 11:01:09.217299
  
  """
  
  from uuid import uuid4
  
  import sqlalchemy as sa
  import sqlmodel
  from alembic import op
  
  from zenml.utils.time_utils import utc_now
  
  # revision identifiers, used by Alembic.
  revision = "a91762e6be36"
  down_revision = "0.50.0"
  branch_labels = None
  depends_on = None
  
  
  def upgrade() -> None:
      """Upgrade database schema and/or data, creating a new revision."""
      # Rename old artifact table to artifact_version
      op.rename_table("artifact", "artifact_version")
  
      # Create new artifact table
      op.create_table(
          "artifact",
          sa.Column("id", sqlmodel.sql.sqltypes.GUID(), nullable=False),
          sa.Column("created", sa.DateTime(), nullable=False),
          sa.Column("updated", sa.DateTime(), nullable=False),
          sa.Column("name", sa.VARCHAR(length=255), nullable=False),
          sa.Column("has_custom_name", sa.Boolean(), nullable=False),
          sa.PrimaryKeyConstraint("id"),
      )
  
      # Add foreign key column to artifact_version, for now as nullable
      with op.batch_alter_table("artifact_version", schema=None) as batch_op:
          batch_op.add_column(
              sa.Column(
                  "artifact_id", sqlmodel.sql.sqltypes.GUID(), nullable=True
              )
          )
  
      # Find all unique artifact names and create a new artifact entry for each
      conn = op.get_bind()
      meta = sa.MetaData()
      meta.reflect(only=("artifact", "artifact_version"), bind=op.get_bind())
      artifacts = sa.Table("artifact", meta)
      artifact_versions = sa.Table("artifact_version", meta)
      artifact_names = conn.execute(
          sa.select(
              artifact_versions.c.name, artifact_versions.c.has_custom_name
          ).distinct()
      ).all()
      unique_artifact_names = {
          name: has_custom_name for name, has_custom_name in artifact_names
      }
      now = utc_now()
      for name, has_custom_name in unique_artifact_names.items():
          conn.execute(
              artifacts.insert().values(
                  id=uuid4().hex,
                  created=now,
                  updated=now,
                  name=name,
                  has_custom_name=has_custom_name,
              )
          )
  
      # Set artifact_id column in artifact_version
      conn.execute(
          artifact_versions.update().values(
              artifact_id=sa.select(artifacts.c.id)
              .where(artifacts.c.name == artifact_versions.c.name)
              .scalar_subquery()
          )
      )
  
      # Make foreign key column non-nullable and create constraint
      with op.batch_alter_table("artifact_version", schema=None) as batch_op:
          batch_op.alter_column(
              "artifact_id", nullable=False, existing_type=sa.CHAR(length=32)
          )
          batch_op.create_foreign_key(
              "fk_artifact_version_artifact_id_artifact",
              "artifact",
              ["artifact_id"],
              ["id"],
              ondelete="CASCADE",
          )
  
      # Drop old columns from artifact_version
      with op.batch_alter_table("artifact_version", schema=None) as batch_op:
          batch_op.drop_column("name")
          batch_op.drop_column("has_custom_name")
  
      # Rename constraints `fk_artifact_...` to `fk_artifact_version_...`
      with op.batch_alter_table("artifact_version", schema=None) as batch_op:
          batch_op.drop_constraint(
              "fk_artifact_artifact_store_id_stack_component", type_="foreignkey"
          )
          batch_op.drop_constraint(
              "fk_artifact_user_id_user", type_="foreignkey"
          )
          batch_op.drop_constraint(
              "fk_artifact_workspace_id_workspace", type_="foreignkey"
          )
      with op.batch_alter_table("artifact_version", schema=None) as batch_op:
          batch_op.create_foreign_key(
              "fk_artifact_version_artifact_store_id_stack_component",
              "stack_component",
              ["artifact_store_id"],
              ["id"],
              ondelete="SET NULL",
          )
          batch_op.create_foreign_key(
              "fk_artifact_version_user_id_user",
              "user",
              ["user_id"],
              ["id"],
              ondelete="SET NULL",
          )
          batch_op.create_foreign_key(
              "fk_artifact_version_workspace_id_workspace",
              "workspace",
              ["workspace_id"],
              ["id"],
              ondelete="CASCADE",
          )
  
      # Rename constraints of other tables
      with op.batch_alter_table(
          "artifact_visualization", schema=None
      ) as batch_op:
          batch_op.drop_constraint(
              "fk_artifact_visualization_artifact_id_artifact",
              type_="foreignkey",
          )
          batch_op.alter_column(
              "artifact_id",
              new_column_name="artifact_version_id",
              existing_type=sa.CHAR(length=32),
          )
      with op.batch_alter_table(
          "artifact_visualization", schema=None
      ) as batch_op:
          batch_op.create_foreign_key(
              "fk_artifact_visualization_artifact_version_id_artifact_version",
              "artifact_version",
              ["artifact_version_id"],
              ["id"],
              ondelete="CASCADE",
          )
  
      with op.batch_alter_table(
          "model_versions_artifacts", schema=None
      ) as batch_op:
          batch_op.drop_constraint(
              "fk_model_versions_artifacts_artifact_id_artifact",
              type_="foreignkey",
          )
          batch_op.alter_column(
              "artifact_id",
              new_column_name="artifact_version_id",
              existing_type=sa.CHAR(length=32),
          )
      with op.batch_alter_table(
          "model_versions_artifacts", schema=None
      ) as batch_op:
          batch_op.create_foreign_key(
              "fk_model_versions_artifacts_artifact_version_id_artifact_version",
              "artifact_version",
              ["artifact_version_id"],
              ["id"],
              ondelete="CASCADE",
          )
  
      with op.batch_alter_table("run_metadata", schema=None) as batch_op:
          batch_op.drop_constraint(
              "fk_run_metadata_artifact_id_artifact", type_="foreignkey"
          )
          batch_op.alter_column(
              "artifact_id",
              new_column_name="artifact_version_id",
              existing_type=sa.CHAR(length=32),
          )
      with op.batch_alter_table("run_metadata", schema=None) as batch_op:
          batch_op.create_foreign_key(
              "fk_run_metadata_artifact_version_id_artifact_version",
              "artifact_version",
              ["artifact_version_id"],
              ["id"],
              ondelete="CASCADE",
          )
  
      with op.batch_alter_table(
          "step_run_input_artifact", schema=None
      ) as batch_op:
          batch_op.drop_constraint(
              "fk_step_run_input_artifact_artifact_id_artifact",
              type_="foreignkey",
          )
          batch_op.create_foreign_key(
              "fk_step_run_input_artifact_artifact_id_artifact_version",
              "artifact_version",
              ["artifact_id"],
              ["id"],
              ondelete="CASCADE",
          )
  
      with op.batch_alter_table(
          "step_run_output_artifact", schema=None
      ) as batch_op:
          batch_op.drop_constraint(
              "fk_step_run_output_artifact_artifact_id_artifact",
              type_="foreignkey",
          )
          batch_op.create_foreign_key(
              "fk_step_run_output_artifact_artifact_id_artifact_version",
              "artifact_version",
              ["artifact_id"],
              ["id"],
              ondelete="CASCADE",
          )
  
  
  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.")
  