module PgConn
  class Error < StandardError; end # Move to exception.rb
  class PsqlError < Error; end

  class RdbmsMethods
    attr_reader :conn

    def initialize(conn)
      @conn = conn
      # TODO: Check if conn is a superuser connection
    end

    # Return true if the database exists
    def exist?(database)
      conn.exist? %(
          select 1
          from pg_database
          where datname = '#{database}'
      )
    end

    # Create a new database
    def create(database, owner: ENV['USER'], template: "template1") 
      owner_clause = owner ? "owner = \"#{owner}\"" : nil
      template_clause = template ? "template = \"#{template}\"" : nil
      stmt = ["create database \"#{database}\"", owner_clause, template_clause].compact.join(" ")
      conn.execute stmt # Note we're using #execute instead of #exec because
                        # create database can't run within a transaction
    end

    # Drop a database
    def drop(database, force: false)
      conn.execute "drop database if exists \"#{database}\"#{force ? ' with (force)' : ''}"
      true
    end

    # List databases in the RDBMS
    def list(all: false, exclude: [])
      exclude += POSTGRES_DATABASES if !all
      exclude_sql_list = "'" + exclude.join("', '") + "'"
      exclude_clause = exclude.empty? ? nil : "where datname not in (#{exclude_sql_list})"
      stmt = ["select datname from pg_database", exclude_clause].compact.join(" ")
      conn.values stmt
    end

    # Return the owner of a given database
    def owner(database)
      conn.value %(
        select r.rolname
        from (values ('#{database}')) as v (database)
        left join pg_database d on d.datname = v.database
        left join pg_roles r on r.oid = d.datdba
      )
    end

    # Return list of users currently logged in to the given database or to any
    # database if database is nil
    #
    # FIXME: There is a possible race-condition here where some process (eg.
    # auto-vacuum) is logged in to the database but has a nil username.  The
    # easy fix is to have 'usename is not null' but it would be nice to know
    # what exactly is triggering this problem
    #
    def users(database)
      database_clause = database ? "datname = '#{database}'" : nil
      query = ["select usename from pg_stat_activity", database_clause].compact.join(" where ")
      conn.values query
    end

    # Hollow-out a database by removing all schemas in the database. The public
    # schema is recreated afterwards unless if :public is false. Uses the current
    # database if @database is nil
    #
    # Note that the database can have active users logged in while the database
    # is emptied
    #
    def empty!(database = nil, public: true, exclude: [])
      local = !database.nil?
      begin
        conn = local ? PgConn.new(database) : self.conn
        schemas = 
            conn
              .values("select nspname from pg_namespace where nspowner != 10 or nspname = 'public'")
              .select { |schema| !exclude.include?(schema) }
              .join(", ")
        conn.exec "drop schema #{schemas} cascade"

        # FIXME SECURITY Why grant 'create' to public?
        conn.exec %(
            create schema public authorization postgres;
            grant usage, create on schema public to public 
        ) if public
      ensure
        conn&.terminate if local
      end
    end

    # Fast copy using templates. Note that no user may be logged in to the
    # source database for this to work
    def copy(from_database, to_database, owner: ENV['USER'])
      create(to_database, owner: owner, template: from_database)
    end

    def load(database, file, role: ENV['USER'], gzip: nil)
      command_opt = role ? "-c \"set role #{role}\";\n" : nil
      if gzip
        pipe_cmd = file ? "gunzip --to-stdout #{file} |" : "gunzip --to-stdout |"
        file_opt = nil
      else
        pipe_cmd = nil
        file_opt = file ? "-f #{file}" : nil
      end
      cmd = [pipe_cmd, "psql -v ON_ERROR_STOP=1", command_opt, file_opt, database].compact.join(" ")
      stdout, stderr, status = Open3.capture3(cmd)
      status == 0 or raise PsqlError.new(stderr)
    end

    def save(database, file, data: true, schema: true, gzip: nil)
      data_opt = data ? nil : "--schema-only"
      schema_opt = schema ? nil : "--data-only"
      if gzip
        pipe_cmd = file ? "| gzip >#{file}" : "| gzip"
        file_opt = nil
      else
        pipe_cmd = nil
        file_opt = file ? "-f #{file}" : nil
      end
      cmd = ["pg_dump --no-owner", data_opt, schema_opt, file_opt, database, pipe_cmd].compact.join(" ")
      stdout, stderr, status = Open3.capture3(cmd)
      status == 0 or raise PsqlError.new(stderr)
    end

  private
    POSTGRES_DATABASES = %w(template0 template1 postgres)
  end
end