def prepare_databases_for_metrics()

in postgresql_metrics/prepare_db.py [0:0]


def prepare_databases_for_metrics(conf):
    """Tries first to connect to localhost database as default user,
    which works if the local user is setup as local Postgres superuser.
    If this fails, queries for Postgres superuser credentials.
    """
    metrics_user = conf['postgres']['user']
    metrics_user_password = conf['postgres']['password']
    LOG.info("prepare databases for metrics user '{}'", metrics_user)

    db_names = []
    if 'databases' in conf['postgres']:
        db_names = conf['postgres']['databases']
    elif 'database' in conf['postgres']:
        db_names = [conf['postgres']['database']]

    for db_name in db_names:
        LOG.info("connecting to database '{}' as super user", db_name)
        db_connection = connect_as_super_user(db_name, conf)

        if check_if_database_is_slave(db_connection):
            LOG.info("database is a slave, run prepare-db on master")
            break

        if not check_if_role_exists(db_connection, metrics_user):
            create_role_with_login(db_connection, metrics_user, metrics_user_password)
        else:
            LOG.info("role already exists: {}", metrics_user)

        if not check_if_role_has_db_privilege(db_connection, metrics_user, db_name, 'connect'):
            LOG.info("grant connect privilege to user '{}' for database: {}",
                     metrics_user, db_name)
            with db_connection.cursor() as c:
                c.execute("GRANT CONNECT ON database " + db_name + " TO " + metrics_user)
        else:
            LOG.info("role '{}' already has connect privilege to database: {}",
                     metrics_user, db_name)

        if not check_if_replication_stats_view_exists(db_connection):
            create_replication_stats_view(db_connection)
        else:
            LOG.info("replication stats view already exists")

        if not check_if_role_has_table_privilege(db_connection, metrics_user,
                                                 REPLICATION_STATS_VIEW, 'select'):
            LOG.info("grant select privilege to user '{}' for relation: {}",
                     metrics_user, REPLICATION_STATS_VIEW)
            with db_connection.cursor() as c:
                c.execute("GRANT SELECT ON " + REPLICATION_STATS_VIEW + " TO " + metrics_user)
        else:
            LOG.info("role '{}' already has select privilege to relation: {}",
                     metrics_user, REPLICATION_STATS_VIEW)

        if not check_if_pgstattuples_extension_exists(db_connection):
            create_pgstattuples_extension(db_connection)
        else:
            LOG.info("pgstattuples extension already exists")

        if not check_if_role_has_function_privilege(db_connection, metrics_user,
                                                    PGSTATTUPLES_FUNC, 'execute'):
            LOG.info("grant execute privilege to user '{}' for function: {}",
                     metrics_user, PGSTATTUPLES_FUNC)
            with db_connection.cursor() as c:
                c.execute("GRANT EXECUTE ON FUNCTION " + PGSTATTUPLES_FUNC + " TO "
                          + metrics_user)
        else:
            LOG.info("role '{}' already has execute privilege to function: {}",
                     metrics_user, PGSTATTUPLES_FUNC)

        if db_connection.server_version >= PGVERSION_WAL_RECEIVER:
            if not check_if_incoming_replication_status_view_exists(db_connection):
                create_incoming_replication_status_view(db_connection)
            else:
                LOG.info("incoming replication status view already exists")

            if not check_if_role_has_table_privilege(db_connection, metrics_user,
                                                     INCOMING_REPLICATION_STATS_VIEW, 'select'):
                LOG.info("grant select privilege to user '{}' for relation: {}",
                         metrics_user, INCOMING_REPLICATION_STATS_VIEW)
                with db_connection.cursor() as c:
                    c.execute(sql.SQL("GRANT SELECT ON {} TO {}").format(
                        sql.Identifier(INCOMING_REPLICATION_STATS_VIEW),
                        sql.Identifier(metrics_user)))
            else:
                LOG.info("role '{}' already has select privilege to relation: {}",
                         metrics_user, REPLICATION_STATS_VIEW)
        else:
            LOG.info("skipping setup for incoming replication view, requires Postgres version >= %s",
                     PGVERSION_WAL_RECEIVER)

        LOG.info("database '{}' prepared for metrics user: {}", db_name, metrics_user)