class Prodder::PG

Constants

ACL_GRANT

From pg_dump

ACL_REVOKE
DEFAULT_PRIVILEGES
PGDumpError
SET_OBJECT_OWNERSHIP

Attributes

credentials[R]

Public Class Methods

new(credentials = {}) click to toggle source
# File lib/prodder/pg.rb, line 10
def initialize(credentials = {})
  @credentials = credentials
end

Public Instance Methods

create_db(db_name, sql = nil) click to toggle source
# File lib/prodder/pg.rb, line 31
def create_db(db_name, sql = nil)
  run ['createdb', db_name]
  psql db_name, sql if sql
end
create_role(role, opts = []) click to toggle source
# File lib/prodder/pg.rb, line 14
def create_role(role, opts = [])
  arguments = [
    'createuser',
    '--no-password',
    '--no-superuser',
    '--no-createrole',
    '--no-createdb'
  ]

  arguments.push *opts, role
  run arguments
end
drop_db(db_name) click to toggle source
# File lib/prodder/pg.rb, line 36
def drop_db(db_name)
  run ['dropdb', db_name]
end
drop_role(role) click to toggle source
# File lib/prodder/pg.rb, line 27
def drop_role(role)
  run ['dropuser', role]
end
dump_db_access_control(db_name, user_list, options) click to toggle source
# File lib/prodder/pg.rb, line 129
def dump_db_access_control(db_name, user_list, options)
  perm_out_sql = ""
  arguments = [
    '--schema-only',
    '--host', credentials['host'],
    '--username', credentials['user']
  ]

  if options[:exclude_schemas].respond_to? :map
    arguments.concat options[:exclude_schemas].map { |schema| ['--exclude-schema', schema] }.flatten
  end

  if options[:exclude_tables].respond_to? :map
    arguments.concat options[:exclude_tables].map { |table| ['--exclude-table', table] }.flatten
  end

  arguments.push db_name

  white_list = options[:included_users] || []
  irrelevant_login_roles = irrelevant_login_roles(db_name, white_list).map { |user| user['rolname'] }

  run ['pg_dump', *arguments] do |out, err, success|
    out.each_line do |line|
      if line.match(ACL_GRANT)           ||
         line.match(ACL_REVOKE)          ||
         line.match(DEFAULT_PRIVILEGES)  ||
         line.match(SET_OBJECT_OWNERSHIP)

        unless irrelevant_login_roles.include?(line.match(/ (\S*);$/)[1])
          perm_out_sql << line
          user_list << (line.match(/ (\S*);$/)[1]).gsub(/"/, '')
        end
      end
    end
    raise PGDumpError.new(err) if !success
  end

  user_list.uniq!
  perm_out_sql
end
dump_permissions(db_name, filename, options = {}) click to toggle source
# File lib/prodder/pg.rb, line 105
def dump_permissions(db_name, filename, options = {})
  perm_out_sql = ""
  user_list = []

  perm_out_sql << dump_db_access_control(db_name, user_list, options)
  perm_out_sql.prepend pg_dumpall db_name, user_list, options

  perm_out_sql.prepend(alter_role_function)
  perm_out_sql.prepend(create_role_function)
  perm_out_sql.prepend(granted_by_function)
  perm_out_sql.prepend("SET client_min_messages TO WARNING;\n")
  perm_out_sql << drop_role_create_function
  perm_out_sql << drop_role_alter_function
  perm_out_sql << drop_granted_by_function

  File.open(filename, 'w') { |f| f.write perm_out_sql }
end
dump_settings(db_name, filename) click to toggle source
# File lib/prodder/pg.rb, line 44
    def dump_settings(db_name, filename)
      db_settings = pg_conn(db_name) { |conn| conn.exec(<<-SQL).map { |setting| setting['config'] } }
        select unnest(setconfig) as config
        from pg_catalog.pg_db_role_setting
        join pg_database on pg_database.oid = setdatabase
        -- 0 = default, for all users
        where setrole = 0
        and datname = '#{db_name}'
      SQL

      File.open(filename, 'w') do |f|
        db_settings.each do |setting|
          # wipe out all spaces
          setting.gsub!(/\s+/, '')

          # if the setting is empty, ignore it
          unless setting.empty?
            # else, drop carriage returns/new lines
            setting.chomp!
            # and append an empty string if the setting was being assigned a value of nothing
            setting += "''" if setting.match(/=$/)
            # using the magic of psql variables through :DBNAME
            f.puts "ALTER DATABASE :DBNAME SET #{setting};"
          end
        end
      end
    end
dump_structure(db_name, filename, options = {}) click to toggle source
# File lib/prodder/pg.rb, line 72
def dump_structure(db_name, filename, options = {})
  arguments = [
    '--schema-only',
    '--no-privileges',
    '--no-owner',
    '--host', credentials['host'],
    '--username', credentials['user']
  ]

  if options[:exclude_schemas].respond_to? :map
    arguments.concat options[:exclude_schemas].map { |schema| ['--exclude-schema', schema] }.flatten
  end

  if options[:exclude_tables].respond_to? :map
    arguments.concat options[:exclude_tables].map { |table| ['--exclude-table', table] }.flatten
  end

  pg_dump filename, arguments.push(db_name)
end
dump_tables(db_name, tables, filename) click to toggle source
# File lib/prodder/pg.rb, line 92
def dump_tables(db_name, tables, filename)
  pg_dump filename, [
    '--data-only',
    '--no-privileges',
    '--no-owner',
    '--disable-triggers',
    '--host',     credentials['host'],
    '--username', credentials['user'],
    *tables.map { |table| ['--table', table] }.flatten,
    db_name
  ]
end
psql(db_name, sql) click to toggle source
# File lib/prodder/pg.rb, line 40
def psql(db_name, sql)
  run ['psql', db_name], sql
end

Private Instance Methods

alter_role_function() click to toggle source
# File lib/prodder/pg.rb, line 343
    def alter_role_function
      <<-SQL
        CREATE OR REPLACE FUNCTION public.alter_role(rolename VARCHAR, sql TEXT)
         RETURNS TEXT
         AS
         $alter_role$
         DECLARE
           r RECORD;
           compensating_sql TEXT := '**!!**ALTER ROLE ';
         BEGIN
           EXECUTE 'SELECT rolsuper, rolinherit,
                           rolcreaterole, rolcreatedb,
                           rolcanlogin, rolconnlimit,
                           rolvaliduntil, rolreplication,
                           pg_catalog.shobj_description(oid, $1) as rolcomment
                    FROM pg_roles
                    WHERE rolname = $2'
           INTO r
           USING 'pg_authid', rolename;
           compensating_sql := compensating_sql || '"' || rolename || '" WITH';
           IF r.rolsuper THEN
             compensating_sql := compensating_sql || ' SUPERUSER';
           ELSE
             compensating_sql := compensating_sql || ' NOSUPERUSER';
           END IF;
           IF r.rolinherit THEN
             compensating_sql := compensating_sql || ' INHERIT';
           ELSE
             compensating_sql := compensating_sql || ' NOINHERIT';
           END IF;
           IF r.rolcreaterole THEN
             compensating_sql := compensating_sql || ' CREATEROLE';
           ELSE
             compensating_sql := compensating_sql || ' NOCREATEROLE';
           END IF;
           IF r.rolcreatedb THEN
             compensating_sql := compensating_sql || ' CREATEDB';
           ELSE
             compensating_sql := compensating_sql || ' NOCREATEDB';
           END IF;
           IF r.rolcanlogin THEN
             compensating_sql := compensating_sql || ' LOGIN';
           ELSE
             compensating_sql := compensating_sql || ' NOLOGIN';
           END IF;
           IF r.rolreplication THEN
             compensating_sql := compensating_sql || ' REPLICATION';
           ELSE
             compensating_sql := compensating_sql || ' NOREPLICATION';
           END IF;
           IF r.rolconnlimit <> -1 THEN
             compensating_sql := compensating_sql || ' CONNECTION LIMIT ' || r.rolconnlimit;
           END IF;
           IF r.rolvaliduntil IS NOT NULL THEN
             compensating_sql := compensating_sql || ' VALID UNTIL ' || r.rolvaliduntil;
           END IF;
           compensating_sql := compensating_sql || ';\n';
           IF r.rolcomment IS NOT NULL THEN
             compensating_sql := compensating_sql || 'COMMENT ON ROLE "' || rolename || '" IS ' || r.rolcomment || ';\n';
           END IF;
           compensating_sql := compensating_sql || '**!!**';
           EXECUTE sql;
           RETURN compensating_sql;
         END;
         $alter_role$
         LANGUAGE PLPGSQL;
      SQL
    end
create_role_function() click to toggle source
# File lib/prodder/pg.rb, line 454
    def create_role_function
      <<-SQL
        \n
        CREATE OR REPLACE FUNCTION public.create_role_if_not_exists(rolename VARCHAR)
        RETURNS TEXT
        AS
        $create_role_if_not_exists$
        DECLARE
        BEGIN
          IF NOT EXISTS (
              SELECT *
              FROM   pg_catalog.pg_roles
              WHERE  rolname = rolename) THEN
            EXECUTE 'CREATE ROLE ' || quote_ident(rolename) || ' ;';
            RETURN '**!!**DROP ROLE ''' || rolename || ''';**!!**';
          ELSE
            RETURN FALSE;
          END IF;
        END;
        $create_role_if_not_exists$
        LANGUAGE PLPGSQL;
        \n
      SQL
    end
drop_granted_by_function() click to toggle source
# File lib/prodder/pg.rb, line 446
    def drop_granted_by_function
      <<-SQL
        \n
        DROP FUNCTION public.granted_by(VARCHAR, TEXT);
        \n
      SQL
    end
drop_role_alter_function() click to toggle source
# File lib/prodder/pg.rb, line 412
    def drop_role_alter_function
      <<-SQL
        \n
        DROP FUNCTION public.alter_role(VARCHAR, TEXT);
        \n
      SQL
    end
drop_role_create_function() click to toggle source
# File lib/prodder/pg.rb, line 479
    def drop_role_create_function
      <<-SQL
        \n
        DROP FUNCTION public.create_role_if_not_exists(VARCHAR);
        \n
      SQL
    end
granted_by_function() click to toggle source
# File lib/prodder/pg.rb, line 420
    def granted_by_function
      <<-SQL
        \n
        CREATE OR REPLACE FUNCTION public.granted_by(rolename VARCHAR, sql TEXT)
        RETURNS BOOLEAN
        AS
        $role_exists$
        DECLARE
        BEGIN
          IF EXISTS (
              SELECT *
              FROM   pg_catalog.pg_roles
              WHERE  rolname = rolename) THEN
            EXECUTE sql;
            RETURN TRUE;
          ELSE
            RAISE NOTICE 'Rolename % does not exist, cannot set granted by', rolename;
            RETURN FALSE;
          END IF;
        END;
        $role_exists$
        LANGUAGE PLPGSQL;
        \n
      SQL
    end
irrelevant_login_roles(db_name, white_list) click to toggle source
# File lib/prodder/pg.rb, line 185
def irrelevant_login_roles(db_name, white_list)
  white_list ||= []
  login_role_list = pg_conn(db_name) do |conn|
    conn.exec('SELECT oid, rolname FROM pg_roles WHERE rolcanlogin AND NOT rolsuper').map do |role|
      {'oid' => role['oid'], 'rolname' => role['rolname'] }
    end
  end

  login_role_list.reject! { |user| white_list.include?(user['rolname']) }
  login_role_list
end
pg_conn(db_name) { |conn| ... } click to toggle source
# File lib/prodder/pg.rb, line 172
def pg_conn(db_name)
  conn = ::PG.connect(
    dbname: db_name,
    host: credentials['host'],
    user: credentials['user'],
    password: credentials['password']
  )

  res = yield(conn)
  conn.close
  res
end
pg_dump(filename, cmd) click to toggle source
# File lib/prodder/pg.rb, line 213
def pg_dump(filename, cmd)
  run ['pg_dump', *cmd] do |out, err, success|
    raise PGDumpError.new(err) if !success
    File.open(filename, 'w') { |f| f.write out }
  end
end
pg_dumpall(db_name, user_list, options) click to toggle source
# File lib/prodder/pg.rb, line 220
def pg_dumpall(db_name, user_list, options)
  white_list = options[:included_users] || []
  irrelevant_login_roles = irrelevant_login_roles(db_name, white_list).map { |user| user['oid'] }

  roles_and_memberships = pg_conn(db_name) { |conn| conn.exec smart_pgdumpall(user_list, irrelevant_login_roles) }

  rolcreate_sql, rolalter_sql, rolgrant_sql = "", "", ""
  created_roles = Set.new

  roles_and_memberships.each do |role|
    unless created_roles.member? role['rolname']
      created_roles << role['rolname']
      tmp_sql = ""
      rolcreate_sql << "SELECT * FROM create_role_if_not_exists('#{role['rolname']}');\n"
      tmp_sql << "ALTER ROLE \"#{role['rolname']}\" WITH"

      [
        ['rolsuper', 'SUPERUSER'],
        ['rolinherit', 'INHERIT'],
        ['rolcreaterole', 'CREATEROLE'],
        ['rolcreatedb', 'CREATEDB'],
        ['rolcanlogin', 'LOGIN'],
        ['rolreplication', 'REPLICATION']
      ].each do |key, modifier|
        tmp_sql << if role[key].eql? 't'
          " #{modifier}"
        else
          modifier.prepend ' NO'
        end
      end

      tmp_sql << " CONNECTION LIMIT #{role['rolconnlimit']}" unless role['rolconnlimit'].eql?("-1")
      tmp_sql << " VALID UNTIL '#{role['rolvaliduntil']}'" unless role['rolvaliduntil'].nil?
      tmp_sql << ";\n"
      tmp_sql << "COMMENT ON ROLE \"#{role['rolname']}\" IS '#{role['rolcomment']}';\n" unless role['rolcomment'].nil?
      rolalter_sql << "SELECT * FROM alter_role('#{role['rolname']}', $$#{tmp_sql}$$);\n"
    end

    unless role['member'].nil?
      tmp_sql = ""
      rolgrant_sql << "GRANT \"#{role['roleid']}\" TO \"#{role['member']}\""
      rolgrant_sql << " WITH ADMIN OPTION" if role['admin_option'].eql?('t')
      rolgrant_sql << ";\n"
      tmp_sql << "GRANT \"#{role['roleid']}\" TO \"#{role['member']}\""
      tmp_sql << " WITH ADMIN OPTION" if role['admin_option'].eql?('t')
      tmp_sql << " GRANTED BY #{role['grantor']}" if role['grantor'].eql?('t')
      tmp_sql << ";"
      rolgrant_sql << "SELECT * FROM granted_by('#{role['grantor']}', $$#{tmp_sql}$$);\n"
    end
  end

  rolcreate_sql << rolalter_sql << rolgrant_sql
end
quote(value) click to toggle source
# File lib/prodder/pg.rb, line 330
def quote(value)
  "'#{quote_string(value.to_s)}'"
end
quote_string(s) click to toggle source
# File lib/prodder/pg.rb, line 326
def quote_string(s)
  s.gsub(/\\/, '\&\&').gsub(/'/, "''")
end
replace_bind_variable(value) click to toggle source
# File lib/prodder/pg.rb, line 314
def replace_bind_variable(value)
  if value.respond_to?(:map)
    if value.respond_to?(:empty?) && value.empty?
      quote(nil)
    else
      value.map { |v| quote(v) }.join(',')
    end
  else
    quote(value)
  end
end
replace_bind_variables(statement, values) click to toggle source
# File lib/prodder/pg.rb, line 334
def replace_bind_variables(statement, values)
  values.each do |value|
    statement.sub!(/\?/) do
      replace_bind_variable(value)
    end
  end
  statement
end
run(cmd, stdin_data = nil, &block) click to toggle source
# File lib/prodder/pg.rb, line 197
def run(cmd, stdin_data = nil, &block)
  # TODO use a tmp .pgpass file instead of $PGPASSWORD
  env = { 'PGPASSWORD' => credentials['password'] }
  Open3.popen3(env, *cmd) do |stdin, out, err, thr|
    if stdin_data
      stdin.write stdin_data
      stdin.close
    end

    out, err = out.read, err.read
    puts err if err
    block.call(out, err, thr.value.success?) if block
    out
  end
end
smart_pgdumpall(user_list, irrelevant_login_roles) click to toggle source
# File lib/prodder/pg.rb, line 274
    def smart_pgdumpall(user_list, irrelevant_login_roles)
      irrelevant_login_roles << -1 if irrelevant_login_roles.respond_to?(:empty?) && irrelevant_login_roles.empty?
      replace_bind_variables(<<-SQL, [user_list, irrelevant_login_roles])
        WITH RECURSIVE memberships(roleid, member, admin_option, grantor) AS (
          SELECT ur.oid AS roleid,
                 NULL::oid AS member,
                 NULL::boolean AS admin_option,
                 NULL::oid AS grantor
          FROM pg_roles ur
          WHERE ur.rolname IN (?)
          UNION
          SELECT COALESCE(a.roleid, r.oid) AS roleid,
                 a.member AS member,
                 a.admin_option AS admin_option,
                 a.grantor AS grantor
          FROM pg_auth_members a
          FULL OUTER JOIN pg_roles r ON FALSE
          JOIN memberships
            ON (memberships.roleid = a.member)
            OR (memberships.roleid = r.oid OR memberships.member = r.oid)
            OR (memberships.roleid = a.roleid AND COALESCE(memberships.member, 0::oid) <> a.member AND a.member NOT IN(?))
        )
        SELECT DISTINCT ON(ur.rolname, um.rolname)
               ur.rolname AS roleid,
               um.rolname AS member,
               memberships.admin_option,
               ug.rolname AS grantor,
               ur.rolname, ur.rolsuper, ur.rolinherit,
               ur.rolcreaterole, ur.rolcreatedb,
               ur.rolcanlogin, ur.rolconnlimit,
               ur.rolvaliduntil, ur.rolreplication,
               pg_catalog.shobj_description(memberships.roleid, 'pg_authid') as rolcomment
        FROM memberships
        LEFT JOIN pg_roles ur on ur.oid = memberships.roleid
        LEFT JOIN pg_roles um on um.oid = memberships.member
        LEFT JOIN pg_roles ug on ug.oid = memberships.grantor
        ORDER BY 1,2 NULLS LAST;
      SQL
    end