content: # Copyright Contributors to the Amundsen project.
  # SPDX-License-Identifier: Apache-2.0
  
  import contextlib
  import logging
  import os
  from typing import Dict
  
  import amundsen_rds
  from alembic import command, script
  from alembic.config import Config
  from alembic.runtime.migration import MigrationContext
  from amundsen_rds.models.base import Base
  from sqlalchemy import create_engine
  from sqlalchemy.orm import Session, sessionmaker
  
  LOGGER = logging.getLogger(__name__)
  
  
  class RDSClient:
      """
      Client used by mysql proxy and rds commands
      """
      def __init__(self, sql_alchemy_url: str, client_kwargs: Dict = dict()) -> None:
          self.sql_alchemy_url = sql_alchemy_url
          self.engine = create_engine(sql_alchemy_url, **client_kwargs)
          self.session_factory = sessionmaker(bind=self.engine)
  
      def init_db(self) -> None:
          """
          Initialize a relational database
          """
          LOGGER.info('Initializing the metadata database')
  
          config = self._get_alembic_config()
          config.set_main_option('sqlalchemy.url', self.sql_alchemy_url.replace('%', '%%'))
  
          command.upgrade(config, 'heads')
  
      def reset_db(self) -> None:
          """
          Reset schema
          """
          LOGGER.info('Resetting the metadata database')
  
          self.drop_models()
          self.init_db()
  
      def drop_models(self) -> None:
          """
          Drop rds models
          """
          LOGGER.info('Dropping rds models')
  
          Base.metadata.drop_all(self.engine)
  
          with self.engine.connect() as conn:
              migration_ctx = MigrationContext.configure(conn)
              version = migration_ctx._version
              if version.exists(conn):
                  version.drop(conn)
  
      def validate_schema_version(self) -> bool:
          """
          Check if the schema from the given connection is the latest one
          """
          LOGGER.info('Validating schema version.')
  
          config = self._get_alembic_config()
          script_directory = script.ScriptDirectory.from_config(config)
  
          with self.engine.connect() as conn:
              context = MigrationContext.configure(conn)
              current_version = context.get_current_revision()
              current_head = script_directory.get_current_head()
  
              LOGGER.info(f'current_version: {current_version}, current_head: {current_head}')
  
          return current_version == current_head
  
      @staticmethod
      def _get_alembic_config() -> Config:
          """
          Get alembic config
          """
          amundsen_rds_dir = os.path.dirname(os.path.abspath(amundsen_rds.__file__))
          migration_path = os.path.join(amundsen_rds_dir, 'migrations')
          config = Config(os.path.join(amundsen_rds_dir, 'alembic.ini'))
          config.set_main_option('script_location', migration_path.replace('%', '%%'))
  
          return config
  
      @contextlib.contextmanager
      def create_session(self) -> Session:
          """
          Create a SQLAlchemy session
          """
          session = self.session_factory()
          try:
              yield session
          except Exception as e:
              session.rollback()
              raise e
          finally:
              session.close()
  