content: """Add artifact unique constraints [c22561cbb3a9].
  
  Revision ID: c22561cbb3a9
  Revises: 0.68.1
  Create Date: 2024-10-17 16:41:25.053677
  
  """
  
  from collections import defaultdict
  from typing import Dict, Set
  
  import sqlalchemy as sa
  from alembic import op
  
  # revision identifiers, used by Alembic.
  revision = "c22561cbb3a9"
  down_revision = "0.68.1"
  branch_labels = None
  depends_on = None
  
  
  def resolve_duplicate_versions() -> None:
      """Resolve duplicate artifact versions."""
      connection = op.get_bind()
  
      meta = sa.MetaData()
      meta.reflect(
          bind=op.get_bind(),
          only=("artifact_version",),
      )
  
      artifact_version_table = sa.Table("artifact_version", meta)
      query = sa.select(
          artifact_version_table.c.id,
          artifact_version_table.c.artifact_id,
          artifact_version_table.c.version,
      )
  
      versions_per_artifact: Dict[str, Set[str]] = defaultdict(set)
  
      for id, artifact_id, version in connection.execute(query).fetchall():
          versions = versions_per_artifact[artifact_id]
          if version in versions:
              for suffix_length in range(4, len(id)):
                  new_version = f"{version}-{id[:suffix_length]}"
                  if new_version not in versions:
                      version = new_version
                      break
  
              connection.execute(
                  sa.update(artifact_version_table)
                  .where(artifact_version_table.c.id == id)
                  .values(version=version)
              )
  
          versions.add(version)
  
  
  def upgrade() -> None:
      """Upgrade database schema and/or data, creating a new revision."""
      # ### commands auto generated by Alembic - please adjust! ###
      resolve_duplicate_versions()
  
      with op.batch_alter_table("artifact", schema=None) as batch_op:
          batch_op.create_unique_constraint("unique_artifact_name", ["name"])
  
      with op.batch_alter_table("artifact_version", schema=None) as batch_op:
          batch_op.create_unique_constraint(
              "unique_version_for_artifact_id", ["version", "artifact_id"]
          )
  
      # ### end Alembic commands ###
  
  
  def downgrade() -> None:
      """Downgrade database schema and/or data back to the previous revision."""
      # ### commands auto generated by Alembic - please adjust! ###
      with op.batch_alter_table("artifact_version", schema=None) as batch_op:
          batch_op.drop_constraint(
              "unique_version_for_artifact_id", type_="unique"
          )
  
      with op.batch_alter_table("artifact", schema=None) as batch_op:
          batch_op.drop_constraint("unique_artifact_name", type_="unique")
  
      # ### end Alembic commands ###
  