File Coverage

blib/lib/Test/PostgreSQL.pm

Criterion	Covered	Total	%
statement	85	340	25.0
branch	17	220	7.7
condition	1	44	2.2
subroutine	20	39	51.2
pod	7	10	70.0
total	130	653	19.9

package Test::PostgreSQL;
use 5.014;
use strict;
use warnings;
use Moo;
use Types::Standard -all;
use Function::Parameters qw(:strict);
use Try::Tiny;
use DBI;
use File::Spec;
use File::Temp;
use File::Which;
use POSIX qw(SIGQUIT SIGKILL WNOHANG getuid setuid);
use User::pwent;

our $VERSION = '1.28';
our $errstr;

# Deprecate use of %Defaults as we want to remove this package global
use Tie::Hash::Method;
tie our %Defaults, 'Tie::Hash::Method', FETCH => sub {
    my $msg = "\nWARNING: using \$Test::PostgreSQL::Defaults is DEPRECATED.";
    if ( $_[1] =~ /^(initdb|postmaster)_args$/ ) {
        $msg .= " Use Test::PostgreSQL->new( extra_$_[1] => ... ) instead.";
    }
    warn $msg;
    return $_[0]->base_hash->{ $_[1] };
  };

%Defaults = (
    auto_start      => 2,
    initdb_args     => '-U postgres -A trust',
    postmaster_args => '-h 127.0.0.1 -F',
);

has dbname => (
    is => 'ro',
    isa => Str,
    default => 'test',
);

has dbowner => (
    is => 'ro',
    isa => Str,
    default => 'postgres',
);

has host => (
    is => 'ro',
    isa => Str,
    default => '127.0.0.1',
);

# Various paths that Postgres gets installed under, sometimes with a version on the end,
# in which case take the highest version. We append /bin/ and so forth to the path later.
# *Note that these are used only if the program isn't already in the path!*
has search_paths => (
  is => "ro",
  isa => ArrayRef,
  builder => "_search_paths",
);

method _search_paths() {
    my @base_paths = (
        # popular installation dir?
        qw(/usr/local/pgsql),
        # ubuntu (maybe debian as well, find the newest version)
        (sort { $b cmp $a } grep { -d $_ } glob "/usr/lib/postgresql/*"),
        # macport
        (sort { $b cmp $a } grep { -d $_ } glob "/opt/local/lib/postgresql*"),
        # Postgresapp.com
        (sort { $b cmp $a } grep { -d $_ } glob "/Applications/Postgres.app/Contents/Versions/*"),
        # BSDs end up with it in /usr/local/bin which doesn't appear to be in the path sometimes:
        "/usr/local",
    );

    # This environment variable is used to override the default, so it gets
    # prefixed to the start of the search paths.
    if (defined $ENV{POSTGRES_HOME}) {
      return [$ENV{POSTGRES_HOME}, @base_paths];
    }
    return \@base_paths;
}

# We attempt to use this port first, and will increment from there.
# The final port ends up in the ->port attribute.
has base_port => (
  is => "ro",
  isa => Int,
  default => 15432,
);

has auto_start => (
  is => "ro",
  default => 2,
);

has base_dir => (
  is => "rw",
  default => sub {
    File::Temp->newdir(
        'pgtest.XXXXX',
        CLEANUP => $ENV{TEST_POSTGRESQL_PRESERVE} ? undef : 1,
        EXLOCK  => 0,
        TMPDIR  => 1
    );
  },
  coerce => fun ($newval) {
    # Ensure base_dir is absolute; usually only the case if the user set it.
    # Avoid munging objects such as File::Temp
    ref $newval ? $newval : File::Spec->rel2abs($newval);
  },
);

has socket_dir => (
  is => "ro",
  isa => Str,
  lazy => 1,
  default => method () { File::Spec->catdir( $self->base_dir, 'tmp' ) },
);

has initdb => (
  is => "ro",
  isa => Str,
  lazy => 1,
  default => method () { $self->_find_program('initdb') || die $errstr },
);

has initdb_args => (
  is => "lazy",
  isa => Str,
);

method _build_initdb_args() {
    return '-U '. $self->dbowner . ' -A trust ' . $self->extra_initdb_args;
}

has extra_initdb_args => (
  is      => "ro",
  isa     => Str,
  default => "",
);

has unix_socket => (
  is  => "ro",
  isa     => Bool,
  default => 0,
);

has pg_version => (
    is => 'ro',
    isa => Str,
    lazy => 1,
    predicate => 1,
    builder => "_pg_version_builder",
);

method _pg_version_builder() {
    my $ver_cmd = join ' ', (
        $self->postmaster,
        '--version'
    );

    my ($ver) = qx{$ver_cmd} =~ /(\d+(?:\.\d+)?)/;

    return $ver;
}

has pg_ctl => (
  is => "ro",
  isa => Maybe[Str],
  lazy => 1,
  builder => "_pg_ctl_builder",
);

method _pg_ctl_builder() {
  my $prog = $self->_find_program('pg_ctl');
  if ( $prog ) {
      # we only use pg_ctl if Pg version is >= 9
      my $ret = qx/"$prog" --version/;
      if ( $ret =~ /(\d+)(?:\.|devel|beta)/ && $1 >= 9 ) {
          return $prog;
      }
      warn "pg_ctl version earlier than 9";
      return;
  }
  return;
}

has pg_config => (
    is => 'ro',
    isa => Str,
);

has psql => (
    is => 'ro',
    isa => Str,
    lazy => 1,
    default => method () { $self->_find_program('psql') || die $errstr },
);

has psql_args => (
    is => 'lazy',
    isa => Str,
);

method _build_psql_args() {
    return '-U ' . $self->dbowner . ' -d ' . $self->dbname . ' -h '.
        ($self->unix_socket ? $self->socket_dir : '127.0.0.1') .
        ' -p ' . $self->port
        . $self->extra_psql_args;
}

has extra_psql_args => (
    is => 'ro',
    isa => Str,
    default => '',
);

has run_psql_args => (
    is => 'ro',
    isa => Str,
    lazy => 1,
    builder => "_build_run_psql_args",
);

method _build_run_psql_args() {
    my @args = (
        '-1', # Single transaction
        '-X', # Ignore .psqlrc
        '-q', # Quiet
        '-v ON_ERROR_STOP=1', # Stop on first error
    );

    # Echo errors, available in psql 9.5+
    push @args, '-b' if $self->pg_version >= 9.5;

    return join ' ', @args;
}

has seed_scripts => (
    is => 'ro',
    isa => ArrayRef[Str],
    default => sub { [] },
);

has pid => (
  is => "rw",
  isa => Maybe[Int],
);

has port => (
  is => "rw",
  isa => Maybe[Int],
);

has uid => (
  is => "rw",
  isa => Maybe[Int],
);

# Are we running as root? (Typical when run inside Docker containers)
has is_root => (
  is => "ro",
  isa => Bool,
  default => sub { getuid == 0 }
);

has postmaster => (
  is => "rw",
  isa => Str,
  lazy => 1,
  default => method () {
    $self->_find_program("postgres")
    || $self->_find_program("postmaster")
    || die $errstr
  },
);

has postmaster_args => (
  is => "lazy",
  isa => Str,
);

method _build_postmaster_args() {
    return "-h ".
        ($self->unix_socket ? "''" : "127.0.0.1") .
        " -F " . $self->extra_postmaster_args;
}

has extra_postmaster_args => (
    is      => "ro",
    isa     => Str,
    default => "",
);

has _owner_pid => (
  is => "ro",
  isa => Int,
  default => sub { $$ },
);

method BUILD($) {
    # Ensure we have one or the other ways of starting Postgres:
    try { $self->pg_ctl or $self->postmaster } catch { die $_ };

    if (defined $self->uid and $self->uid == 0) {
        die "uid() must be set to a non-root user id.";
    }

    if (not defined($self->uid) and $self->is_root) {
        my $ent = getpwnam("nobody");
        unless (defined $ent) {
            die "user nobody does not exist, use uid() to specify a non-root user.";
        }
        unless ($ent->uid > 0) {
            die "user nobody has uid 0; confused and exiting. use uid() to specify a non-root user.";
        }
        $self->uid($ent->uid);
    }

    # Ensure base dir is writable by our target uid, if we were running as root
    chown $self->uid, -1, $self->base_dir
        if defined $self->uid;

    if ($self->auto_start) {
        $self->setup
            if $self->auto_start >= 2;
        $self->start;
    }
}

method DEMOLISH($in_global_destruction) {
    local $?;
    if (defined $self->pid && $self->_owner_pid == $$) {
      $self->stop
    }
    return;
}

sub dsn {
    my %args = shift->_default_args(@_);

    return 'DBI:Pg:' . join(';', map { "$_=$args{$_}" } sort keys %args);
}

sub _default_args {
    my ($self, %args) = @_;
    # If we're doing socket-only (i.e., not listening on localhost),
    # then provide the path to the socket
    if ($self->{unix_socket}) {
        $args{host} //= $self->socket_dir;
    } else {
        $args{host} ||= $self->host;
    }

    $args{port} ||= $self->port;
    $args{user} ||= $self->dbowner;
    $args{dbname} ||= $self->dbname;
    return %args;
}

sub uri {
    my $self = shift;
    my %args = $self->_default_args(@_);

    return sprintf('postgresql://%s@%s:%d/%s', @args{qw/user host port dbname/});
}

method start() {
    if (defined $self->pid) {
      warn "Apparently already started on " . $self->pid . "; not restarting.";
      return;
    }

    # If the user specified a port, try only that port:
    if ($self->port) {
        $self->_try_start($self->port);
    }
    else {
        $self->_find_port_and_launch;
    }

    # create "test" database
    $self->_create_test_database($self->dbname);
}

# This whole method was mostly cargo-culted from the earlier test-postgresql;
# It could probably be made more sane.
method _find_port_and_launch() {
  my $tries = 10;
  my $port = $self->base_port;
  srand(); # Re-seed the RNG in case the caller forked the process
  # try by incrementing port number until PostgreSQL starts
  while (1) {
    my $good = try {
      $self->_try_start($port);
      1;
    }
    catch {
      # warn "Postgres failed to start on port $port\n";
      unless ($tries--) {
        die "Failed to start postgres after trying 10 potential ports: $_";
      }
      undef;
    };
    return if $good;
    # Increment port by a random number to avoid clashes with other Test::Postgresql processes
    # Keep in mind that this increment is going to be made up to 10 times, so avoid exceeding 64k
    $port += int(rand(500)) + 1;
  }
}

method _try_start($port) {
    my $logfile = File::Spec->catfile($self->base_dir, 'postgres.log');

    if ( $self->pg_ctl ) {
        my @cmd = (
            $self->pg_ctl,
            'start', '-w', '-s', '-D',
            File::Spec->catdir( $self->base_dir, 'data' ),
            '-l', $logfile, '-o',
            join( ' ',
                $self->postmaster_args, '-p',
                $port,                  '-k',
                $self->socket_dir)
        );
        $self->setuid_cmd(\@cmd, 1);

        my $pid_path = File::Spec->catfile( $self->base_dir, 'data', 'postmaster.pid' );

        open( my $pidfh, '<', $pid_path )
          or die "Failed to open $pid_path: $!";

        # Note that the file contains several lines; we only want the PID from the first.
        my $pid = <$pidfh>;
        chomp $pid;
        $self->pid($pid);
        close $pidfh;

        $self->port($port);
    }
    else {
        # old style - open log and fork
        open my $logfh, '>>', $logfile
            or die "failed to create log file: $logfile: $!";
        my $pid = fork;
        die "fork(2) failed:$!"
            unless defined $pid;
        if ($pid == 0) {
            open STDOUT, '>>&', $logfh
                or die "dup(2) failed:$!";
            open STDERR, '>>&', $logfh
                or die "dup(2) failed:$!";
            chdir $self->base_dir
                or die "failed to chdir to:" . $self->base_dir . ":$!";
            if (defined $self->uid) {
                setuid($self->uid) or die "setuid failed: $!";
            }
            my $cmd = join(
                ' ',
                $self->postmaster,
                $self->postmaster_args,
                '-p', $port,
                '-D', File::Spec->catdir($self->base_dir, 'data'),
                '-k', $self->socket_dir,
            );
            exec($cmd);
            die "failed to launch postmaster:$?";
        }
        close $logfh;
        # wait until server becomes ready (or dies)
        for (my $i = 0; $i < 100; $i++) {
            open $logfh, '<', $logfile
            or die "failed to create log file: $logfile: $!";
            my $lines = do { join '', <$logfh> };
            close $logfh;
            last
                if $lines =~ /is ready to accept connections/;
            if (waitpid($pid, WNOHANG) > 0) {
                # failed
                die "Failed to start Postgres: $lines\n";
            }
            sleep 1;
        }
        # PostgreSQL is ready
        $self->pid($pid);
        $self->port($port);
    }
    return;
}

method stop($sig = SIGQUIT) {
    if ( $self->pg_ctl && defined $self->base_dir ) {
        my @cmd = (
            $self->pg_ctl, 'stop', '-s', '-D',
            File::Spec->catdir( $self->base_dir, 'data' ),
            '-m', 'fast'
        );
        $self->setuid_cmd(\@cmd);
    }
    else {
        # old style or $self->base_dir File::Temp obj already DESTROYed
        return unless defined $self->pid;

        kill $sig, $self->pid;
        my $timeout = 10;
        while ($timeout > 0 and waitpid($self->pid, WNOHANG) == 0) {
            $timeout -= sleep(1);
        }

        if ($timeout <= 0) {
            warn "Pg refused to die gracefully; killing it violently.\n";
            kill SIGKILL, $self->pid;
            $timeout = 5;
            while ($timeout > 0 and waitpid($self->pid, WNOHANG) == 0) {
                $timeout -= sleep(1);
            }
            if ($timeout <= 0) {
                warn "Pg really didn't die.. WTF?\n";
            }
        }
    }
    $self->pid(undef);
    return;
}

method _create_test_database($dbname) {
  my $tries = 5;
  my $dbh;
  while ($tries) {
      $tries -= 1;
      $dbh = DBI->connect($self->dsn(dbname => 'template1'), '', '', {
          PrintError => 0,
          RaiseError => 0
      });
      last if $dbh;

      # waiting for database to start up
      if ($DBI::errstr =~ /the database system is starting up/
          || $DBI::errstr =~ /Connection refused/) {
          sleep(1);
          next;
      }
      die $DBI::errstr;
  }

  die "Connection to the database failed even after 5 tries"
      unless ($dbh);

  if ($dbh->selectrow_arrayref(qq{SELECT COUNT(*) FROM pg_database WHERE datname='$dbname'})->[0] == 0) {
      $dbh->do("CREATE DATABASE $dbname")
          or die $dbh->errstr;
  }

    my $seed_scripts = $self->seed_scripts || [];

    $self->run_psql_scripts(@$seed_scripts)
        if @$seed_scripts;

    return;
}

method setup() {
    # (re)create directory structure
    mkdir $self->base_dir;
    chmod 0755, $self->base_dir
        or die "failed to chmod 0755 dir:" . $self->base_dir . ":$!";
    if ($ENV{USER} && $ENV{USER} eq 'root') {
        chown $self->uid, -1, $self->base_dir
            or die "failed to chown dir:" . $self->base_dir . ":$!";
    }
    my $tmpdir = $self->socket_dir;
    if (mkdir $tmpdir) {
        if ($self->uid) {
            chown $self->uid, -1, $tmpdir
                or die "failed to chown dir:$tmpdir:$!";
        }
    }
    # initdb
    if (! -d File::Spec->catdir($self->base_dir, 'data')) {
        if ( $self->pg_ctl ) {
            my @cmd = (
                $self->pg_ctl,
                'init',
                '-s',
                '-D', File::Spec->catdir($self->base_dir, 'data'),
                '-o',
                $self->initdb_args,
            );
            $self->setuid_cmd(\@cmd);
        }
        else {
            # old style
            pipe my $rfh, my $wfh
                or die "failed to create pipe:$!";
            my $pid = fork;
            die "fork failed:$!"
                unless defined $pid;
            if ($pid == 0) {
                close $rfh;
                open STDOUT, '>&', $wfh
                    or die "dup(2) failed:$!";
                open STDERR, '>&', $wfh
                    or die "dup(2) failed:$!";
                chdir $self->base_dir
                    or die "failed to chdir to:" . $self->base_dir . ":$!";
                if (defined $self->uid) {
                    setuid($self->uid)
                        or die "setuid failed:$!";
                }
                my $cmd = join(
                    ' ',
                    $self->initdb,
                    $self->initdb_args,
                    '-D', File::Spec->catdir($self->base_dir, 'data'),
                );
                exec($cmd);
                die "failed to exec:$cmd:$!";
            }
            close $wfh;
            my $output = '';
            while (my $l = <$rfh>) {
                $output .= $l;
            }
            close $rfh;
            while (waitpid($pid, 0) <= 0) {
            }
            die "*** initdb failed ***\n$output\n"
                if $? != 0;

        }

        my $conf_file
            = File::Spec->catfile($self->base_dir, 'data', 'postgresql.conf');

        if (my $pg_config = $self->pg_config) {
            open my $fh, '>', $conf_file or die "Can't open $conf_file: $!";
            print $fh $pg_config;
            close $fh;
        }
        else {
            # use postgres hard-coded configuration as some packagers mess
            # around with postgresql.conf.sample too much:
            truncate $conf_file, 0;
        }
    }
}

method _find_program($prog) {
    undef $errstr;
    my $path = which $prog;
    return $path if $path;
    for my $sp (@{$self->search_paths}) {
        return "$sp/bin/$prog" if -x "$sp/bin/$prog";
        return "$sp/$prog" if -x "$sp/$prog";
    }
    $errstr = "could not find $prog, please set appropriate PATH or POSTGRES_HOME";
    return;
}

method setuid_cmd($cmd, $suppress_errors = !1) {
  my $pid = fork;
  if ($pid == 0) {
    chdir $self->base_dir;
    if (defined $self->uid) {
      setuid($self->uid) or die "setuid failed: $!";
    }
    close STDERR if $suppress_errors;
    exec(@$cmd) or die "Failed to exec pg_ctl: $!";
  }
  else {
    waitpid($pid, 0);
  }
}

method run_psql(@psql_args) {
    my $cmd = join ' ', (
        $self->psql,

        # Default connection settings
        $self->psql_args,

        # Extra connection settings or something else
        $self->extra_psql_args,

        # run_psql specific arguments
        $self->run_psql_args,

        @psql_args,
    );

    # Usually anything less than WARNING is not really helpful
    # in batch mode. Does it make sense to make this configurable?
    local $ENV{PGOPTIONS} = '--client-min-messages=warning';

    my $psql_out = qx{$cmd 2>&1};

    die "Error executing psql: $psql_out" unless $? == 0;
}

method run_psql_scripts(@script_paths) {
    my @psql_commands;

    # psql 9.6+ supports multiple -c and -f commands invoked at once,
    # older psql does not. Executing psql multiple times breaks single
    # transaction semantics but is unlikely to cause problems in real world.
    if ( $self->pg_version > 9.6 ) {
        push @psql_commands, join ' ', map {; "-f $_" } @script_paths;
    }
    else {
        @psql_commands = map {; "-f $_" } @script_paths;
    }

    $self->run_psql($_) for @psql_commands;
}

1;
__END__