text: # This file is part of Indico.
  # Copyright (C) 2002 - 2022 CERN
  #
  # Indico is free software; you can redistribute it and/or
  # modify it under the terms of the MIT License; see the
  # LICENSE file for more details.
  
  import os
  
  import alembic.command
  import click
  from alembic.script import ScriptDirectory
  from flask import current_app
  from flask_migrate import Migrate, stamp
  from flask_pluginengine import current_plugin
  
  from indico.core.db import db
  from indico.core.db.sqlalchemy.util.management import create_all_tables, get_all_tables
  from indico.core.db.sqlalchemy.util.queries import has_extension
  from indico.core.plugins import plugin_engine
  from indico.util.console import cformat
  
  
  migrate = Migrate(db=db)
  
  
  class PluginScriptDirectory(ScriptDirectory):
      """Like `ScriptDirectory` but lets you override the paths from outside.
  
      This is a pretty ugly hack but alembic doesn't give us a nice way to do it...
      """
  
      dir = None
      versions = None
  
      def __init__(self, *args, **kwargs):
          super().__init__(*args, **kwargs)
          self.dir = PluginScriptDirectory.dir
          # use __dict__ since it's a memoized property
          self.__dict__['_version_locations'] = [current_plugin.alembic_versions_path]
  
      @classmethod
      def from_config(cls, config):
          instance = super().from_config(config)
          instance.dir = PluginScriptDirectory.dir
          instance.__dict__['_version_locations'] = [current_plugin.alembic_versions_path]
          return instance
  
  
  def _require_extensions(*names):
      missing = sorted(name for name in names if not has_extension(db.engine, name))
      if not missing:
          return True
      click.secho(f"Required Postgres extensions missing: {', '.join(missing)}", fg='red')
      click.secho('Create them using these SQL commands (as a Postgres superuser):', fg='yellow')
      for name in missing:
          click.secho(f'  CREATE EXTENSION {name};', bold=True)
      return False
  
  
  def _require_pg_version(version):
      # convert version string such as '9.4.10' to `90410` which is the
      # format used by server_version_num
      # FIXME: this will not work for versions >= 10, since the second segment there is the patch version
      # but once we require a newer postgres, we can ditch the logic here and only use the major version
      # since that will be the only relevant version number
      req_version = sum(segment * 10**(4 - 2*i) for i, segment in enumerate(map(int, version.split('.'))))
      cur_version = db.engine.execute("SELECT current_setting('server_version_num')::int").scalar()
      if cur_version >= req_version:
          return True
      click.secho(f'Postgres version too old; you need at least {version} (or newer)', fg='red')
      return False
  
  
  def _require_encoding(encoding):
      cur_encoding = db.engine.execute("SELECT current_setting('server_encoding')").scalar()
      if cur_encoding == encoding:
          return True
      click.secho(f'Database encoding must be {encoding}; got {cur_encoding}', fg='red')
      click.secho(f'Recreate your database using `createdb -E {encoding} -T template0 ...`', fg='yellow')
      return False
  
  
  def prepare_db(empty=False, root_path=None, verbose=True):
      """Initialize an empty database (create tables, set alembic rev to HEAD)."""
      if not _require_pg_version('9.6'):
          return
      if not _require_encoding('UTF8'):
          return
      if not _require_extensions('unaccent', 'pg_trgm'):
          return
      root_path = root_path or current_app.root_path
      tables = get_all_tables(db)
      if 'alembic_version' not in tables['public']:
          if verbose:
              click.secho('Setting the alembic version to HEAD', fg='green')
          stamp(directory=os.path.join(root_path, 'migrations'), revision='heads')
          PluginScriptDirectory.dir = os.path.join(root_path, 'core', 'plugins', 'alembic')
          alembic.command.ScriptDirectory = PluginScriptDirectory
          plugin_msg = cformat('%{cyan}Setting the alembic version of the %{cyan!}{}%{reset}%{cyan} '
                               'plugin to HEAD%{reset}')
          for plugin in plugin_engine.get_active_plugins().values():
              if not os.path.exists(plugin.alembic_versions_path):
                  continue
              if verbose:
                  print(plugin_msg.format(plugin.name))
              with plugin.plugin_context():
                  stamp(revision='heads')
          # Retrieve the table list again, just in case we created unexpected tables
          tables = get_all_tables(db)
  
      tables['public'] = [t for t in tables['public'] if not t.startswith('alembic_version')]
      if any(tables.values()):
          if verbose:
              click.secho('Your database is not empty!', fg='red')
              click.secho('If you just added a new table/model, create an alembic revision instead!', fg='yellow')
              print()
              print('Tables in your database:')
              for schema, schema_tables in sorted(tables.items()):
                  for t in schema_tables:
                      print(cformat('  * %{cyan}{}%{reset}.%{cyan!}{}%{reset}').format(schema, t))
          return
      create_all_tables(db, verbose=verbose, add_initial_data=(not empty))
  