content: #  Copyright (c) ZenML GmbH 2022. All Rights Reserved.
  #
  #  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:
  #
  #       https://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.
  """ZenML database migration, backup and recovery utilities."""
  
  import json
  import os
  import re
  import shutil
  from typing import (
      Any,
      Callable,
      Dict,
      Generator,
      List,
      Optional,
      cast,
  )
  
  import pymysql
  from pydantic import BaseModel, ConfigDict
  from sqlalchemy import MetaData, func, text
  from sqlalchemy.engine import URL, Engine
  from sqlalchemy.exc import (
      OperationalError,
  )
  from sqlalchemy.schema import CreateIndex, CreateTable
  from sqlmodel import (
      create_engine,
      select,
  )
  
  from zenml.logger import get_logger
  from zenml.utils.json_utils import pydantic_encoder
  
  logger = get_logger(__name__)
  
  
  class MigrationUtils(BaseModel):
      """Utilities for database migration, backup and recovery."""
  
      url: URL
      connect_args: Dict[str, Any]
      engine_args: Dict[str, Any]
  
      _engine: Optional[Engine] = None
      _master_engine: Optional[Engine] = None
  
      def create_engine(self, database: Optional[str] = None) -> Engine:
          """Get the SQLAlchemy engine for a database.
  
          Args:
              database: The name of the database. If not set, a master engine
                  will be returned.
  
          Returns:
              The SQLAlchemy engine.
          """
          url = self.url._replace(database=database)
          return create_engine(
              url=url,
              connect_args=self.connect_args,
              **self.engine_args,
          )
  
      @property
      def engine(self) -> Engine:
          """The SQLAlchemy engine.
  
          Returns:
              The SQLAlchemy engine.
          """
          if self._engine is None:
              self._engine = self.create_engine(database=self.url.database)
          return self._engine
  
      @property
      def master_engine(self) -> Engine:
          """The SQLAlchemy engine for the master database.
  
          Returns:
              The SQLAlchemy engine for the master database.
          """
          if self._master_engine is None:
              self._master_engine = self.create_engine()
          return self._master_engine
  
      @classmethod
      def is_mysql_missing_database_error(cls, error: OperationalError) -> bool:
          """Checks if the given error is due to a missing database.
  
          Args:
              error: The error to check.
  
          Returns:
              If the error because the MySQL database doesn't exist.
          """
          from pymysql.constants.ER import BAD_DB_ERROR
  
          if not isinstance(error.orig, pymysql.err.OperationalError):
              return False
  
          error_code = cast(int, error.orig.args[0])
          return error_code == BAD_DB_ERROR
  
      def database_exists(
          self,
          database: Optional[str] = None,
      ) -> bool:
          """Check if a database exists.
  
          Args:
              database: The name of the database to check. If not set, the
                  database name from the configuration will be used.
  
          Returns:
              Whether the database exists.
  
          Raises:
              OperationalError: If connecting to the database failed.
          """
          database = database or self.url.database
  
          engine = self.create_engine(database=database)
          try:
              engine.connect()
          except OperationalError as e:
              if self.is_mysql_missing_database_error(e):
                  return False
              else:
                  logger.exception(
                      f"Failed to connect to mysql database `{database}`.",
                  )
                  raise
          else:
              return True
  
      def drop_database(
          self,
          database: Optional[str] = None,
      ) -> None:
          """Drops a mysql database.
  
          Args:
              database: The name of the database to drop. If not set, the
                  database name from the configuration will be used.
          """
          database = database or self.url.database
          with self.master_engine.connect() as conn:
              # drop the database if it exists
              logger.info(f"Dropping database '{database}'")
              conn.execute(text(f"DROP DATABASE IF EXISTS `{database}`"))
  
      def create_database(
          self,
          database: Optional[str] = None,
          drop: bool = False,
      ) -> None:
          """Creates a mysql database.
  
          Args:
              database: The name of the database to create. If not set, the
                  database name from the configuration will be used.
              drop: Whether to drop the database if it already exists.
          """
          database = database or self.url.database
          if drop:
              self.drop_database(database=database)
  
          with self.master_engine.connect() as conn:
              logger.info(f"Creating database '{database}'")
              conn.execute(text(f"CREATE DATABASE IF NOT EXISTS `{database}`"))
  
      def backup_database_to_storage(
          self, store_db_info: Callable[[Dict[str, Any]], None]
      ) -> None:
          """Backup the database to a storage location.
  
          Backup the database to an abstract storage location. The storage
          location is specified by a function that is called repeatedly to
          store the database information. The function is called with a single
          argument, which is a dictionary containing either the table schema or
          table data. The dictionary contains the following keys:
  
              * `table`: The name of the table.
              * `create_stmt`: The table creation statement.
              * `data`: A list of rows in the table.
  
          Args:
              store_db_info: The function to call to store the database
                  information.
          """
          metadata = MetaData()
          metadata.reflect(bind=self.engine)
          with self.engine.connect() as conn:
              for table in metadata.sorted_tables:
                  # 1. extract the table creation statements
  
                  create_table_construct = CreateTable(table)
                  create_table_stmt = str(create_table_construct).strip()
                  for column in create_table_construct.columns:
                      # enclosing all column names in backticks. This is because
                      # some column names are reserved keywords in MySQL. For
                      # example, keys and values. So, instead of tracking all
                      # keywords, we just enclose all column names in backticks.
                      # enclose the first word in the column definition in
                      # backticks
                      words = str(column).split()
                      words[0] = f"`{words[0]}`"
                      create_table_stmt = create_table_stmt.replace(
                          f"\n\t{str(column)}", " ".join(words)
                      )
                  # if any double quotes are used for column names, replace them
                  # with backticks
                  create_table_stmt = create_table_stmt.replace('"', "") + ";"
  
                  # enclose all table names in backticks. This is because some
                  # table names are reserved keywords in MySQL (e.g key
                  # and trigger).
                  create_table_stmt = create_table_stmt.replace(
                      f"CREATE TABLE {table.name}",
                      f"CREATE TABLE `{table.name}`",
                  )
                  # do the same for references to other tables
                  # (i.e. foreign key constraints) by replacing REFERENCES <word>
                  # with REFERENCES `<word>`
                  # use a regular expression for this
                  create_table_stmt = re.sub(
                      r"REFERENCES\s+(\w+)",
                      r"REFERENCES `\1`",
                      create_table_stmt,
                  )
  
                  # In SQLAlchemy, the CreateTable statement may not always
                  # include unique constraints explicitly if they are implemented
                  # as unique indexes instead. To make sure we get all unique
                  # constraints, including those implemented as indexes, we
                  # extract the unique constraints from the table schema and add
                  # them to the create table statement.
  
                  # Extract the unique constraints from the table schema
                  index_create_statements = []
                  unique_constraints = []
                  for index in table.indexes:
                      if index.unique:
                          unique_columns = [
                              f"`{column.name}`" for column in index.columns
                          ]
                          unique_constraints.append(
                              f"UNIQUE KEY `{index.name}` ({', '.join(unique_columns)})"
                          )
                      else:
                          if index.name in {
                              fk.name for fk in table.foreign_key_constraints
                          }:
                              # Foreign key indices are already handled by the
                              # table creation statement.
                              continue
  
                          index_create = str(CreateIndex(index)).strip()
                          index_create = index_create.replace(
                              f"CREATE INDEX {index.name}",
                              f"CREATE INDEX `{index.name}`",
                          )
                          index_create = index_create.replace(
                              f"ON {table.name}", f"ON `{table.name}`"
                          )
  
                          for column_name in index.columns.keys():
                              # We need this logic here to avoid the column names
                              # inside the index name
                              index_create = index_create.replace(
                                  f"({column_name}", f"(`{column_name}`"
                              )
                              index_create = index_create.replace(
                                  f"{column_name},", f"`{column_name}`,"
                              )
                              index_create = index_create.replace(
                                  f"{column_name})", f"`{column_name}`)"
                              )
  
                          index_create = index_create.replace('"', "") + ";"
                          index_create_statements.append(index_create)
  
                  # Add the unique constraints to the create table statement
                  if unique_constraints:
                      # Remove the closing parenthesis, semicolon and any
                      # whitespaces at the end of the create table statement
                      create_table_stmt = re.sub(
                          r"\s*\)\s*;\s*$", "", create_table_stmt
                      )
                      create_table_stmt = (
                          create_table_stmt
                          + ", \n\t"
                          + ", \n\t".join(unique_constraints)
                          + "\n);"
                      )
  
                  # Detect self-referential foreign keys from the table schema
                  has_self_referential_foreign_keys = False
                  for fk in table.foreign_keys:
                      # Check if the foreign key points to the same table
                      if fk.column.table == table:
                          has_self_referential_foreign_keys = True
                          break
  
                  # Store the table schema
                  store_db_info(
                      dict(
                          table=table.name,
                          create_stmt=create_table_stmt,
                          self_references=has_self_referential_foreign_keys,
                      )
                  )
  
                  for stmt in index_create_statements:
                      store_db_info(
                          dict(
                              table=table.name,
                              index_create_stmt=stmt,
                          )
                      )
  
                  # 2. extract the table data in batches
                  order_by = [col for col in table.primary_key]
  
                  # Fetch the number of rows in the table
                  row_count = conn.scalar(
                      select(func.count()).select_from(table)
                  )
  
                  # Fetch the data from the table in batches
                  if row_count is not None:
                      batch_size = 100
                      for i in range(0, row_count, batch_size):
                          rows = conn.execute(
                              table.select()
                              .order_by(*order_by)
                              .limit(batch_size)
                              .offset(i)
                          ).fetchall()
  
                          store_db_info(
                              dict(
                                  table=table.name,
                                  data=[row._asdict() for row in rows],
                              ),
                          )
  
      def restore_database_from_storage(
          self, load_db_info: Callable[[], Generator[Dict[str, Any], None, None]]
      ) -> None:
          """Restore the database from a backup storage location.
  
          Restores the database from an abstract storage location. The storage
          location is specified by a function that is called repeatedly to
          load the database information from the external storage chunk by chunk.
          The function must yield a dictionary containing either the table schema
          or table data. The dictionary contains the following keys:
  
              * `table`: The name of the table.
              * `create_stmt`: The table creation statement.
              * `data`: A list of rows in the table.
  
          The function must return `None` when there is no more data to load.
  
          Args:
              load_db_info: The function to call to load the database
                  information.
          """
          # Drop and re-create the primary database
          self.create_database(drop=True)
  
          metadata = MetaData()
  
          with self.engine.begin() as connection:
              # read the DB information one JSON object at a time
              self_references: Dict[str, bool] = {}
              for table_dump in load_db_info():
                  table_name = table_dump["table"]
                  if "create_stmt" in table_dump:
                      # execute the table creation statement
                      connection.execute(text(table_dump["create_stmt"]))
                      # Reload the database metadata after creating the table
                      metadata.reflect(bind=self.engine)
                      self_references[table_name] = table_dump.get(
                          "self_references", False
                      )
  
                  if "index_create_stmt" in table_dump:
                      # execute the index creation statement
                      connection.execute(text(table_dump["index_create_stmt"]))
                      # Reload the database metadata after creating the index
                      metadata.reflect(bind=self.engine)
  
                  if "data" in table_dump:
                      # insert the data into the database
                      table = metadata.tables[table_name]
                      if self_references.get(table_name, False):
                          # If the table has self-referential foreign keys, we
                          # need to disable the foreign key checks before inserting
                          # the rows and re-enable them afterwards. This is because
                          # the rows need to be inserted in the correct order to
                          # satisfy the foreign key constraints and we don't sort
                          # the rows by creation time in the backup.
                          connection.execute(text("SET FOREIGN_KEY_CHECKS = 0"))
  
                      for row in table_dump["data"]:
                          # Convert column values to the correct type
                          for column in table.columns:
                              # Blob columns are stored as binary strings
                              if column.type.python_type is bytes and isinstance(
                                  row[column.name], str
                              ):
                                  # Convert the string to bytes
                                  row[column.name] = bytes(
                                      row[column.name], "utf-8"
                                  )
  
                      # Insert the rows into the table in batches
                      batch_size = 100
                      for i in range(0, len(table_dump["data"]), batch_size):
                          connection.execute(
                              table.insert().values(
                                  table_dump["data"][i : i + batch_size]
                              )
                          )
  
                      if table_dump.get("self_references", False):
                          # Re-enable the foreign key checks after inserting the rows
                          connection.execute(text("SET FOREIGN_KEY_CHECKS = 1"))
  
      def backup_database_to_file(self, dump_file: str) -> None:
          """Backup the database to a file.
  
          This method dumps the entire database into a JSON file. Instead of
          using a SQL dump, we use a proprietary JSON dump because:
  
              * it is (mostly) not dependent on the SQL dialect or database version
              * it is safer with respect to SQL injection attacks
              * it is easier to read and debug
  
          The JSON file contains a list of JSON objects instead of a single JSON
          object, because it allows for buffered reading and writing of the file
          and thus reduces the memory footprint. Each JSON object can contain
          either schema or data information about a single table. For tables with
          a large amount of data, the data is split into multiple JSON objects
          with the first object always containing the schema.
  
          The format of the dump is as depicted in the following example:
  
          ```json
          {
              "table": "table1",
              "create_stmt": "CREATE TABLE table1 (id INTEGER NOT NULL, "
                  "name VARCHAR(255), PRIMARY KEY (id))"
          }
          {
              "table": "table1",
              "data": [
              {
                  "id": 1,
                  "name": "foo"
              },
              {
                  "id": 1,
                  "name": "bar"
              },
              ...
              ]
          }
          {
              "table": "table1",
              "data": [
              {
                  "id": 101,
                  "name": "fee"
              },
              {
                  "id": 102,
                  "name": "bee"
              },
              ...
              ]
          }
          ```
  
          Args:
              dump_file: The path to the dump file.
          """
          # create the directory if it does not exist
          dump_path = os.path.dirname(os.path.abspath(dump_file))
          if not os.path.exists(dump_path):
              os.makedirs(dump_path)
  
          if self.url.drivername == "sqlite":
              # For a sqlite database, we can just make a copy of the database
              # file
              assert self.url.database is not None
              shutil.copyfile(
                  self.url.database,
                  dump_file,
              )
              return
  
          with open(dump_file, "w") as f:
  
              def json_dump(obj: Dict[str, Any]) -> None:
                  """Dump a JSON object to the dump file.
  
                  Args:
                      obj: The JSON object to dump.
                  """
                  # Write the data to the JSON file. Use an encoder that
                  # can handle datetime, Decimal and other types.
                  json.dump(
                      obj,
                      f,
                      indent=4,
                      default=pydantic_encoder,
                  )
                  f.write("\n")
  
              # Call the generic backup method with a function that dumps the
              # JSON objects to the dump file
              self.backup_database_to_storage(json_dump)
  
          logger.debug(f"Database backed up to {dump_file}")
  
      def restore_database_from_file(self, dump_file: str) -> None:
          """Restore the database from a backup dump file.
  
          See the documentation of the `backup_database_to_file` method for
          details on the format of the dump file.
  
          Args:
              dump_file: The path to the dump file.
  
          Raises:
              RuntimeError: If the database cannot be restored successfully.
          """
          if not os.path.exists(dump_file):
              raise RuntimeError(
                  f"Database backup file '{dump_file}' does not "
                  f"exist or is not accessible."
              )
  
          if self.url.drivername == "sqlite":
              # For a sqlite database, we just overwrite the database file
              # with the backup file
              assert self.url.database is not None
              shutil.copyfile(
                  dump_file,
                  self.url.database,
              )
              return
  
          # read the DB dump file one JSON object at a time
          with open(dump_file, "r") as f:
  
              def json_load() -> Generator[Dict[str, Any], None, None]:
                  """Generator that loads the JSON objects in the dump file.
  
                  Yields:
                      The loaded JSON objects.
                  """
                  buffer = ""
                  while True:
                      chunk = f.readline()
                      if not chunk:
                          break
                      buffer += chunk
                      if chunk.rstrip() == "}":
                          yield json.loads(buffer)
                          buffer = ""
  
              # Call the generic restore method with a function that loads the
              # JSON objects from the dump file
              self.restore_database_from_storage(json_load)
  
          logger.info(f"Database successfully restored from '{dump_file}'")
  
      def backup_database_to_memory(self) -> List[Dict[str, Any]]:
          """Backup the database in memory.
  
          Returns:
              The in-memory representation of the database backup.
  
          Raises:
              RuntimeError: If the database cannot be backed up successfully.
          """
          if self.url.drivername == "sqlite":
              # For a sqlite database, this is not supported.
              raise RuntimeError(
                  "In-memory backup is not supported for sqlite databases."
              )
  
          db_dump: List[Dict[str, Any]] = []
  
          def store_in_mem(obj: Dict[str, Any]) -> None:
              """Store a JSON object in the in-memory database backup.
  
              Args:
                  obj: The JSON object to store.
              """
              db_dump.append(obj)
  
          # Call the generic backup method with a function that stores the
          # JSON objects in the in-memory database backup
          self.backup_database_to_storage(store_in_mem)
  
          logger.debug("Database backed up in memory")
  
          return db_dump
  
      def restore_database_from_memory(
          self, db_dump: List[Dict[str, Any]]
      ) -> None:
          """Restore the database from an in-memory backup.
  
          Args:
              db_dump: The in-memory database backup to restore from generated
                  by the `backup_database_to_memory` method.
  
          Raises:
              RuntimeError: If the database cannot be restored successfully.
          """
          if self.url.drivername == "sqlite":
              # For a sqlite database, this is not supported.
              raise RuntimeError(
                  "In-memory backup is not supported for sqlite databases."
              )
  
          def load_from_mem() -> Generator[Dict[str, Any], None, None]:
              """Generator that loads the JSON objects from the in-memory backup.
  
              Yields:
                  The loaded JSON objects.
              """
              for obj in db_dump:
                  yield obj
  
          # Call the generic restore method with a function that loads the
          # JSON objects from the in-memory database backup
          self.restore_database_from_storage(load_from_mem)
  
          logger.info("Database successfully restored from memory")
  
      @classmethod
      def _copy_database(cls, src_engine: Engine, dst_engine: Engine) -> None:
          """Copy the database from one engine to another.
  
          This method assumes that the destination database exists and is empty.
  
          Args:
              src_engine: The source SQLAlchemy engine.
              dst_engine: The destination SQLAlchemy engine.
          """
          src_metadata = MetaData()
          src_metadata.reflect(bind=src_engine)
  
          dst_metadata = MetaData()
          dst_metadata.reflect(bind=dst_engine)
  
          # @event.listens_for(src_metadata, "column_reflect")
          # def generalize_datatypes(inspector, tablename, column_dict):
          #     column_dict["type"] = column_dict["type"].as_generic(allow_nulltype=True)
  
          # Create all tables in the target database
          for table in src_metadata.sorted_tables:
              table.create(bind=dst_engine)
  
          # Refresh target metadata after creating the tables
          dst_metadata.clear()
          dst_metadata.reflect(bind=dst_engine)
  
          # Copy all data from the source database to the destination database
          with src_engine.begin() as src_conn:
              with dst_engine.begin() as dst_conn:
                  for src_table in src_metadata.sorted_tables:
                      dst_table = dst_metadata.tables[src_table.name]
                      insert = dst_table.insert()
  
                      # If the table has a `created` column, we use it to sort
                      # the rows in the table starting with the oldest rows.
                      # This is to ensure that the rows are inserted in the
                      # correct order, since some tables have inner foreign key
                      # constraints.
                      if "created" in src_table.columns:
                          order_by = [src_table.columns["created"]]
                      else:
                          order_by = []
                      if "id" in src_table.columns:
                          # If the table has an `id` column, we also use it to
                          # sort the rows in the table, even if we already use
                          # "created" to sort the rows. We need a unique field to
                          # sort the rows, to break the tie between rows with the
                          # same "created" date, otherwise the same entry might
                          # end up multiple times in subsequent pages.
                          order_by.append(src_table.columns["id"])
  
                      row_count = src_conn.scalar(
                          select(func.count()).select_from(src_table)
                      )
  
                      # Copy rows in batches
                      if row_count is not None:
                          batch_size = 50
                          for i in range(0, row_count, batch_size):
                              rows = src_conn.execute(
                                  src_table.select()
                                  .order_by(*order_by)
                                  .limit(batch_size)
                                  .offset(i)
                              ).fetchall()
  
                              dst_conn.execute(
                                  insert, [row._asdict() for row in rows]
                              )
  
      def backup_database_to_db(self, backup_db_name: str) -> None:
          """Backup the database to a backup database.
  
          Args:
              backup_db_name: Backup database name to backup to.
          """
          # Re-create the backup database
          self.create_database(
              database=backup_db_name,
              drop=True,
          )
  
          backup_engine = self.create_engine(database=backup_db_name)
  
          self._copy_database(self.engine, backup_engine)
  
          logger.debug(
              f"Database backed up to the `{backup_db_name}` backup database."
          )
  
      def restore_database_from_db(self, backup_db_name: str) -> None:
          """Restore the database from the backup database.
  
          Args:
              backup_db_name: Backup database name to restore from.
  
          Raises:
              RuntimeError: If the backup database does not exist.
          """
          if not self.database_exists(database=backup_db_name):
              raise RuntimeError(
                  f"Backup database `{backup_db_name}` does not exist."
              )
  
          backup_engine = self.create_engine(database=backup_db_name)
  
          # Drop and re-create the primary database
          self.create_database(
              drop=True,
          )
  
          self._copy_database(backup_engine, self.engine)
  
          logger.debug(
              f"Database restored from the `{backup_db_name}` backup database."
          )
  
      model_config = ConfigDict(arbitrary_types_allowed=True)
  