content: """Migrate onboarding state [b4fca5241eea].
  
  Revision ID: b4fca5241eea
  Revises: 0.61.0
  Create Date: 2024-06-20 15:01:22.414801
  
  """
  
  import json
  from typing import Dict, List
  
  import sqlalchemy as sa
  from alembic import op
  
  # revision identifiers, used by Alembic.
  revision = "b4fca5241eea"
  down_revision = "0.61.0"
  branch_labels = None
  depends_on = None
  
  
  ONBOARDING_KEY_MAPPING = {
      "connect_zenml": ["device_verified"],
      "run_first_pipeline": ["pipeline_run", "starter_setup_completed"],
      "run_remote_pipeline": [
          "production_setup_completed",
      ],
  }
  
  
  def upgrade() -> None:
      """Upgrade database schema and/or data, creating a new revision."""
      connection = op.get_bind()
  
      meta = sa.MetaData()
      meta.reflect(only=("server_settings",), bind=connection)
  
      server_settings_table = sa.Table("server_settings", meta)
  
      existing_onboarding_state = connection.execute(
          sa.select(server_settings_table.c.onboarding_state)
      ).scalar_one_or_none()
  
      new_state = []
  
      if existing_onboarding_state:
          # There was already an existing onboarding state in the DB
          # -> Migrate to the new server keys
          state = json.loads(existing_onboarding_state)
  
          if isinstance(state, Dict):
              for key in state.keys():
                  if key in ONBOARDING_KEY_MAPPING:
                      new_state.extend(ONBOARDING_KEY_MAPPING[key])
          elif isinstance(state, List):
              # Somehow the state is already converted, probably shouldn't happen
              return
  
      # We now query the DB and complete all onboarding steps that we can detect
      # from the database
      meta = sa.MetaData()
      meta.reflect(
          only=(
              "pipeline_run",
              "stack_component",
              "stack",
              "stack_composition",
              "pipeline_deployment",
          ),
          bind=connection,
      )
  
      pipeline_run_table = sa.Table("pipeline_run", meta)
      stack_component_table = sa.Table("stack_component", meta)
      stack_table = sa.Table("stack", meta)
      stack_composition_table = sa.Table("stack_composition", meta)
      pipeline_deployment_table = sa.Table("pipeline_deployment", meta)
  
      pipeline_run_count = connection.execute(
          sa.select(sa.func.count(pipeline_run_table.c.id))
      ).scalar()
      if pipeline_run_count and pipeline_run_count > 0:
          new_state.extend(ONBOARDING_KEY_MAPPING["run_first_pipeline"])
  
      stack_with_remote_orchestrator_count = connection.execute(
          sa.select(sa.func.count(stack_table.c.id))
          .where(stack_composition_table.c.stack_id == stack_table.c.id)
          .where(
              stack_composition_table.c.component_id
              == stack_component_table.c.id
          )
          .where(
              stack_component_table.c.flavor.not_in(["local", "local_docker"])
          )
          .where(stack_component_table.c.type == "orchestrator")
      ).scalar()
      if (
          stack_with_remote_orchestrator_count
          and stack_with_remote_orchestrator_count > 0
      ):
          new_state.append("stack_with_remote_orchestrator_created")
  
      pipeline_run_with_remote_artifact_store_count = connection.execute(
          sa.select(sa.func.count(pipeline_run_table.c.id))
          .where(
              pipeline_run_table.c.deployment_id
              == pipeline_deployment_table.c.id
          )
          .where(pipeline_deployment_table.c.stack_id == stack_table.c.id)
          .where(stack_composition_table.c.stack_id == stack_table.c.id)
          .where(
              stack_composition_table.c.component_id
              == stack_component_table.c.id
          )
          .where(stack_component_table.c.flavor != "local")
          .where(stack_component_table.c.type == "artifact_store")
      ).scalar()
      if (
          pipeline_run_with_remote_artifact_store_count
          and pipeline_run_with_remote_artifact_store_count > 0
      ):
          new_state.append("production_setup_completed")
  
      pipeline_run_with_remote_orchestrator_count = connection.execute(
          sa.select(sa.func.count(pipeline_run_table.c.id))
          .where(
              pipeline_run_table.c.deployment_id
              == pipeline_deployment_table.c.id
          )
          .where(pipeline_deployment_table.c.stack_id == stack_table.c.id)
          .where(stack_composition_table.c.stack_id == stack_table.c.id)
          .where(
              stack_composition_table.c.component_id
              == stack_component_table.c.id
          )
          .where(
              stack_component_table.c.flavor.not_in(["local", "local_docker"])
          )
          .where(stack_component_table.c.type == "orchestrator")
      ).scalar()
      if (
          pipeline_run_with_remote_orchestrator_count
          and pipeline_run_with_remote_orchestrator_count > 0
      ):
          new_state.append("pipeline_run_with_remote_orchestrator")
          new_state.append("production_setup_completed")
  
      if new_state:
          # If any of the items are finished, we also complete the initial
          # onboarding step which is not explicitly tracked in the database
          new_state.append("device_verified")
  
          # Remove duplicate keys
          new_state = list(set(new_state))
  
      connection.execute(
          sa.update(server_settings_table).values(
              onboarding_state=json.dumps(new_state)
          )
      )
  
  
  def downgrade() -> None:
      """Downgrade database schema and/or data back to the previous revision."""
      # ### commands auto generated by Alembic - please adjust! ###
      pass
      # ### end Alembic commands ###
  