text: # Licensed under the Apache License, Version 2.0 (the "License"); you may
  # not use this file except in compliance with the License. You may obtain
  # a copy of the License at
  #
  # http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
  # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
  # License for the specific language governing permissions and limitations
  # under the License.
  
  """add_availability_zones_table
  
  Revision ID: 1f0bd302c1a6
  Revises: 579c267fbb4d
  Create Date: 2015-07-24 12:09:36.008570
  
  """
  
  # revision identifiers, used by Alembic.
  revision = '1f0bd302c1a6'
  down_revision = '579c267fbb4d'
  
  from alembic import op
  from oslo_utils import timeutils
  from oslo_utils import uuidutils
  from sqlalchemy import Column, DateTime, ForeignKey, String, UniqueConstraint
  
  from manila.db.migrations import utils
  
  
  def collect_existing_az_from_services_table(connection, services_table,
                                              az_table):
      az_name_to_id_mapping = dict()
      existing_az = []
      for service in connection.execute(services_table.select()):
          if service.availability_zone in az_name_to_id_mapping:
              continue
  
          az_id = uuidutils.generate_uuid()
          az_name_to_id_mapping[service.availability_zone] = az_id
          existing_az.append({
              'created_at': timeutils.utcnow(),
              'id': az_id,
              'name': service.availability_zone
          })
  
      op.bulk_insert(az_table, existing_az)
  
      return az_name_to_id_mapping
  
  
  def upgrade():
      connection = op.get_bind()
  
      # Create new AZ table and columns
      availability_zones_table = op.create_table(
          'availability_zones',
          Column('created_at', DateTime),
          Column('updated_at', DateTime),
          Column('deleted_at', DateTime),
          Column('deleted', String(length=36), default='False'),
          Column('id', String(length=36), primary_key=True, nullable=False),
          Column('name', String(length=255)),
          UniqueConstraint('name', 'deleted', name='az_name_uc'),
          mysql_engine='InnoDB',
          mysql_charset='utf8')
  
      for table_name, fk_name in (('services', 'service_az_id_fk'),
                                  ('share_instances', 'si_az_id_fk')):
          op.add_column(
              table_name,
              Column('availability_zone_id', String(36),
                     ForeignKey('availability_zones.id', name=fk_name))
          )
  
      # Collect existing AZs from services table
      services_table = utils.load_table('services', connection)
      az_name_to_id_mapping = collect_existing_az_from_services_table(
          connection, services_table, availability_zones_table)
  
      # Map string AZ names to ID's in target tables
      # pylint: disable=no-value-for-parameter
      set_az_id_in_table = lambda table, id, name: (
          op.execute(
              table.update().where(table.c.availability_zone == name).values(
                  {'availability_zone_id': id})
          )
      )
  
      share_instances_table = utils.load_table('share_instances', connection)
      for name, id in az_name_to_id_mapping.items():
          for table_name in [services_table, share_instances_table]:
              set_az_id_in_table(table_name, id, name)
  
      # Remove old AZ columns from tables
      op.drop_column('services', 'availability_zone')
      op.drop_column('share_instances', 'availability_zone')
  
  
  def downgrade():
      connection = op.get_bind()
  
      # Create old AZ fields
      op.add_column('services', Column('availability_zone', String(length=255)))
      op.add_column('share_instances',
                    Column('availability_zone', String(length=255)))
  
      # Migrate data
      az_table = utils.load_table('availability_zones', connection)
      share_instances_table = utils.load_table('share_instances', connection)
      services_table = utils.load_table('services', connection)
  
      for az in connection.execute(az_table.select()):
          # pylint: disable=no-value-for-parameter
          op.execute(
              share_instances_table.update().where(
                  share_instances_table.c.availability_zone_id == az.id
              ).values({'availability_zone': az.name})
          )
          op.execute(
              services_table.update().where(
                  services_table.c.availability_zone_id == az.id
              ).values({'availability_zone': az.name})
          )
  
      # Remove AZ_id columns and AZ table
      op.drop_constraint('service_az_id_fk', 'services', type_='foreignkey')
      op.drop_column('services', 'availability_zone_id')
      op.drop_constraint('si_az_id_fk', 'share_instances', type_='foreignkey')
      op.drop_column('share_instances', 'availability_zone_id')
      op.drop_table('availability_zones')
  