pg_sample

#!/usr/bin/env perl

our $VERSION = "1.17";

=head1 NAME

pg_sample - extract a small, sample dataset from a larger PostgreSQL
database while maintaining referential integrity.

=head1 SYNOPSIS

pg_sample [ option... ] [ dbname ]

=head1 DESCRIPTION

pg_sample is a utility for exporting a small, sample dataset from a
larger PostgreSQL database. The output and command-line options closely
resemble the pg_dump backup utility (although only the plain-text format
is supported).

The sample database produced includes all tables from the original,
maintains referential integrity, and supports circular dependencies.

To build an actual instance of the sample database, the output of this script
can be piped to the psql utility. For example, assuming we have an existing
PostgreSQL database named "mydb", a sample database could be constructed with:

  createdb sampledb
  pg_sample mydb | psql -v ON_ERROR_STOP=1 sampledb

The "-v ON_ERROR_STOP=1" option is not required but is recommended to catch any
import errors.

=head2 Requirements

* PostgreSQL 8.1 or later

* pg_dump should be in your search path (in order to dump the schema)

* Perl DBI and DBD::Pg (>= 2.0) modules

=head2 Command-line Options

=over

=item I<dbname>

Specifies the database to sample. If not specified, uses the
environment variable PGDATABASE, if defined; otherwise, uses
the username of the user executing the script.

=item B<-a>

=item B<--data-only>

Output only the data, not the schema (data definitions).

=item B<--help>

Output detailed options and exit.

=item B<-E> I<encoding>

=item B<--encoding=>I<encoding>

Use the specified character set encoding. If not specified, uses the
environment variable PGCLIENTENCODING, if defined; otherwise, uses
the encoding of the database.

=item B<-f> I<file>

=item B<--file=>I<file>

Send output to the specified file. If omitted, standard output is used.

=item B<--force>

Drop the sample schema if it exists.

=item B<--keep>

Don't delete the sample schema when the script finishes.

=item B<--limit=>I<limit>

As a numeric value, specifies the default number of rows to copy from
each table (defaults to 100). Note that sample tables may end up with
significantly more rows in order to satisfy foreign key constraints.

If the value is a string, it is interpreted as a pattern/rule pair to
apply to matching tables. Examples:

  # include all rows from the users table
  --limit="users = *"

  # include 1,000 rows from users table
  --limit="users = 1000"

  # include all users where deactivated column is false
  --limit="users = NOT deactivated"

  # include all rows from all tables in the forums schema
  --limit="forums.* = *"

The limit option may be specified multiple times. Multiple pattern/rule
pairs can also be specified as a single comma-separated value. For example:

  # include all rows from the ads table; otherwise default to 300 rows
  --limit="ads=*,*=300"

Rules are applied in order with the first match taking precedence.

=item B<--ordered>
Select from the table records after ordering the records by primary key value
in descending order. --ordered-desc and --ordered-asc are also available to
control whether sort is descending or ascending, respectively.

=item B<--random>

Randomize the rows initially selected from each table. May significantly
increase the running time of the script.

=item B<--sample-schema=>I<schema>

Schema name to use for the sample database (defaults to _pg_sample).

=item B<--schema=>I<schema>

Limit sampling to the specified schema. By default, all schemas are sampled.

=item B<--trace>

Turn on Perl DBI tracing. See the DBI module documentation for details.

=item B<--verbose>

Output status information to standard error.

=back

The following options control the database connection parameters.

=over

=item B<-h> I<host>

=item B<--host>=I<host>

The host name to connect to. Defaults to the PGHOST environment
variable if not specified.

=item B<-p> I<port>

=item B<--port=>I<port>

The database port to connect to. Defaults to the PGPORT environment
variable, if set; otherwise, the default port is used.

=item B<-U> I<username>

=item B<--username=>I<username>

User name to connect as.

=item B<-W> I<password>

=item B<-password=>I<password>

Password to connect with.

=back

=head1 LICENSE

This code is released under the Artistic License. See L<perlartistic>.

=head1 SEE ALSO

L<createdb(1)>, L<pg_dump(1)>, L<psql(1)>

=head1 AUTHOR

Maurice Aubrey <maurice.aubrey@gmail.com>

=cut

# Algorithm: we create a special sample schema in the original database
# (named _pg_sample by default).
#
# For each table in the original database, a similarly named table is created
# in the sample schema, along with a subset of rows (100 rows maximum, by
# default).
#
# Foreign key constraints are then checked, and any rows needed to satisfy them
# are copied into the sample tables.
#
# Finally, the database schema and sample data are output, renaming the sample
# table names to reflect the original database.

use strict;
use warnings;
use Carp;
use Getopt::Long qw/ GetOptions :config no_ignore_case /;
use DBI;
use DBD::Pg 2.0.0;

our $MAX_IDENTIFIER_LENGTH = 63;

$SIG{TERM} = $SIG{INT} = $SIG{QUIT} = $SIG{HUP} = sub {
  my $signal = shift;
  die "Received signal '$signal'; cleaning up and terminating.\n";
};

my $exit_code = 0;
$SIG{__DIE__} = sub {
  $exit_code = 1 unless $^S; # unless in eval
};

BEGIN {
  # Encapsulate a fully-qualified Postgresql table name
  package Table;

  use strict;
  use Carp;
  use overload '""' => 'quoted'; # interpolation results in quoted version

  *dbh = \&main::connect_db;

  sub new {
    my $class = shift;

    $class = ref($class) || $class;
    return bless {
      schema => shift,
      table => shift,
      columns => shift,
    };
  }

  sub quoted {
    my $self = shift;

    return $self->dbh->quote_identifier(
      undef,
      $self->{schema},
      $self->{table},
    );
  }

  # only used for command-line matching to keep it simple
  sub unquoted {
    my $self = shift;

    return join '.', $self->{schema}, $self->{table};
  }

  sub schema {
    my $self = shift;

    $self->{schema} = shift if @_;
    return $self->{schema};
  }

  sub columns {
    my $self = shift;
    return $self->{columns};
  }

  sub quote_column {
    my $self = shift;
    my $column = shift;
    my $alias = shift;
    if (defined($alias)) {
      return "$alias." . $self->dbh->quote_identifier($column);
    } else {
      return $self->dbh->quote_identifier($column);
    }
  }

  sub columns_quoted {
    my $self = shift;
    my $alias = shift;
    return join(
      ', ',
      map { $self->quote_column($_, $alias) } @{$self->{columns}}
    );
  }

  sub table {
    my $self = shift;

    $self->{table} = shift if @_;
    return $self->{table};
  }

  sub DESTROY {}

  sub AUTOLOAD {
    my $self = shift;

    our $AUTOLOAD;
    $AUTOLOAD =~ s/.*:://g;

    if ($AUTOLOAD =~ /^quoted_(.+)$/) {
      my $method = $1;
      return $self->dbh->quote_identifier($self->$method);
    }

    croak "Can't locate object method \"$AUTOLOAD\" via package " .
          __PACKAGE__;
  }
}

sub list {
  map { ref $_ eq 'ARRAY' ? @$_ : $_ } @_
}

my %opt; # closure; all functions have access to options

{
  my $dbh; # cached db handle

  sub connect_db {
    return $dbh if $dbh;

    my $dsn = join ';',
      "dbi:Pg:dbname=$opt{db_name}",
      $opt{db_host} ? "host=$opt{db_host}" : (),
      $opt{db_port} ? "port=$opt{db_port}" : (),
    ;

    $dbh = DBI->connect(
      $dsn,
      $opt{db_user},
      $opt{db_pass},
      {
        RaiseError       => 1,
        AutoCommit       => 1,
        # FetchHashKeyName => 'NAME_lc',
        PrintError       => 0,
        HandleError      => sub { confess( shift ) },
      },
    ) or croak "db connection failed!";

    $dbh->trace(1) if defined $opt{trace};

    return $dbh;
  }
}

# Retrieve PostgreSQL server version.
# Expect something like:
#   PostgreSQL 8.4.8 on x86_64-pc-linux-gnu, compiled by GCC gcc-4.6.real (Ubuntu/Linaro 4.6.0-7ubuntu1) 4.6.1, 64-bit
sub pg_version {
  my ($ver) = connect_db()->selectrow_array("SELECT version()")
    or croak "request to select postgresql server version failed";
  $ver =~ /([\d.]+)/ or croak "unable to parse postgresql version from '$ver'";

  return version->declare($1);
}

{
  my $maxlen;

  sub max_identifier_length () {
    $maxlen //= do {
      my $len = eval {
        connect_db()->selectrow_array(
          "SELECT current_setting('max_identifier_length')"
        );
      };
      if ($@) {
        warn "Unable to determine max identifier length: $@\n";
        warn "Defaulting to max identifier length $MAX_IDENTIFIER_LENGTH\n";
        $len= $MAX_IDENTIFIER_LENGTH;
      }
      $len;
    };
  }
}


# Return a copy of a hash with all keys lowercased.
sub lower_keys($) {
  my $hash = shift or return;

  ref $hash eq 'HASH' or croak "not a hash reference";

  my %lower = map { lc($_) => $hash->{$_} } keys %$hash;
  return \%lower;
}

# See "Identifiers and Key Words" section 4.1.1
# http://www.postgresql.org/docs/8.2/static/sql-syntax-lexical.html
sub unquote_identifier (@) {
  my @unquoted;
  while (@_) {
    my $val = shift @_;
    if ($val =~ /^\s*"(.*)"\s*$/) { # looks quoted?
      $val = $1;
      $val =~ s/(["\\])\1/$1/g; # undo doubled chars
    }
    push @unquoted, $val;
  }

  return wantarray ? @unquoted : $unquoted[0];
}

# See "String Constants" section 4.1.2.1
# http://www.postgresql.org/docs/8.2/static/sql-syntax-lexical.html
sub quote_constant (@) {
  my @quoted;
  while (@_) {
    my $val = shift @_;
    $val =~ s/'/''/g;
    push @quoted, "'$val'";
  }

  return wantarray ? @quoted: $quoted[0];
}

# Encode the actual schema and table name into a new table
# name that lives under our sample schema. e.g., a table like
# users.details (schema users, table details) would be converted
# to something like _pg_sample.users_details (depending on the
# value of --schema).
#
# @param Table instance
{
  my $seq = 0;
  my %cache;

  sub sample_table ($) {
    my $table = shift;

    my $key = join '_', $table->schema || 'public', $table->table;

    my $sample_table = $cache{ $key } //= do {
      my $name;
      do {
        # make sure sample table names are unique and not too long
        $name = sprintf '%04x__%s', ++$seq, $key;
        $name = substr $name, 0, max_identifier_length;
      } until !exists $cache{ $name };
      $name;
    };

    return Table->new($opt{sample_schema}, $sample_table, $table->columns);
  }
}


# Find the first candidate key for a given table (preferring primary
# key if available)
sub find_candidate_key {
  my $table = shift;

  my $query = qq{
      SELECT n.nspname AS table_schema
             ,c2.relname as table_name
             ,c.relname AS index_name
             ,ar.attname AS column_name
        FROM pg_catalog.pg_class c
             LEFT JOIN pg_catalog.pg_index i ON i.indexrelid = c.oid
             LEFT JOIN pg_catalog.pg_class c2 ON i.indrelid = c2.oid
             LEFT JOIN pg_catalog.pg_namespace n ON n.oid = c2.relnamespace
             LEFT JOIN pg_catalog.pg_attribute ar ON (
               ar.attrelid = c2.oid
               AND ar.attnum = ANY(i.indkey)
               AND NOT ar.attisdropped
             )
       WHERE c.relkind = 'i'
             AND n.nspname = ?
             AND c2.relname = ?
             AND (indisprimary OR indisunique)
             -- exclude partial and functional indexes since we
             -- don't know how to express them from perl
             AND indexprs IS NULL AND indpred IS NULL
    ORDER BY i.indisprimary DESC -- prefer primary key
             ,index_name
             ,column_name
  };
  my $dbh = connect_db;
  my $sth = $dbh->prepare($query);
  $sth->execute($table->schema, $table->table);

  my $last_index_name;
  my @cols;
  while (my $row = $sth->fetchrow_hashref) {
    last if $row->{index_name} ne ($last_index_name //= $row->{index_name});
    push @cols, $row->{column_name};
  }

  return wantarray ? @cols : \@cols;
}


sub notice (@) {
  return unless $opt{verbose};
  print STDERR join "", @_;
}


%opt = (
  db_host       => '',
  db_port       => '',
  keep          => 0,
  ordered       => 0,
  random        => 0,
  schema        => undef,
  sample_schema => '_pg_sample',
  verbose       => 0,
);

GetOptions(\%opt,
  "data-only|data_only|a",
  "db_name|db-name=s",
  "db_user|db-user|db_username|db-username|username|U=s",
  "db_pass|db-pass|db_password|db-password|password|W=s",
  "db_host|db-host|h|host=s",
  "db_port|db-port|p|port=i",
  "encoding|E=s",
  "file|f=s",
  "force",
  "help|?|usage",
  "keep",
  "limit=s@",
  "no-privileges|no-acl|x",
  "ordered",
  "ordered_desc|ordered-desc",
  "ordered_asc|ordered-asc",
  "random",
  "sample_schema=s",
  "schema=s",
  "trace",
  "verbose|v",
  "version|V",
);

if ($opt{version}) {
  print "$VERSION\n";
  exit 0;
}

if ($opt{help}) {
  require Pod::Usage;
  Pod::Usage::pod2usage(-verbose => 2);
  exit 0;
}

$opt{ordered} = $opt{ordered_desc} ? 'DESC'
              : $opt{ordered_asc}  ? 'ASC'
              : $opt{ordered}      ? 'DESC'
              : undef;
if ($opt{random} && $opt{ordered}) {
  print("Error: --random and --ordered are mutually exclusive");
  exit 1;
}

@ARGV or die "\nUsage: $0 [ option... ] [ dbname ]\n\n\t" .
             "$0 --help for detailed options\n";

push @{ $opt{limit} }, ".* = 100 "; # append default limit rule

$opt{db_user}  ||= $ENV{PGUSER} || scalar getpwuid($<);
$opt{db_name}  ||= shift() || $ENV{PGDATABASE} || $opt{db_user};
$opt{db_host}  ||= $ENV{PGHOST} if defined $ENV{PGHOST};
$opt{db_port}  ||= $ENV{PGPORT} if defined $ENV{PGPORT};
$opt{encoding} ||= $ENV{PGCLIENTENCODING} if defined $ENV{PGCLIENTENCODING};

my $dbh = connect_db(%opt) or croak "unable to connect to database";

my $pg_version = pg_version;

if ($opt{sample_schema} eq 'public') {
  die "Error: refusing to use 'public' schema for sampling.\n";
}

my ($schema_oid) = $dbh->selectrow_array(qq{
  SELECT oid
    FROM pg_catalog.pg_namespace
   WHERE nspname = ?
}, undef, $opt{sample_schema});
if ($schema_oid && !$opt{force}) {
  die "Error: schema '$opt{sample_schema}' already exists. " .
      "Use --force option to overwrite.\n";
}

$dbh->do(qq{ SET client_min_messages = warning }); # suppress notice messages
if ($opt{force}) {
  notice "Dropping sample schema $opt{sample_schema}\n";
  $dbh->do(qq{ DROP SCHEMA IF EXISTS $opt{sample_schema} CASCADE });
}

if ($opt{file}) {
  open STDOUT, '>', $opt{file} or croak "unable to redirect stdout: $!";
}

my ($server_encoding) = $dbh->selectrow_array(qq{ SHOW server_encoding });
notice "Server encoding is $server_encoding\n";

$opt{encoding} ||= $server_encoding;
notice "Client encoding is $opt{encoding}\n";
binmode STDOUT, ":encoding($opt{encoding})";

unless ($opt{'data-only'}) {
  notice "Exporting schema\n";

  local $ENV{PGUSER} = $opt{db_user};
  local $ENV{PGPASSWORD} = $opt{db_pass};
  local $ENV{PGDATABASE} = $opt{db_name};
  local $ENV{PGHOST} = $opt{db_host};
  local $ENV{PGPORT} = $opt{db_port};
  local $ENV{PGCLIENTENCODING} = $opt{encoding};

  my @cmd = ('pg_dump', '--schema-only');
  push @cmd, '--no-privileges' if $opt{'no-privileges'};
  push @cmd, "--schema=$opt{schema}" if $opt{schema};
  system(@cmd) == 0 or croak "command '@cmd' failed: $?";
}

# If running PostgreSQL 9.1 or later, use UNLOGGED tables
my $unlogged = $pg_version >= version->declare('9.1') ? 'UNLOGGED' : '';

notice "Creating sample schema $opt{sample_schema}\n";
$dbh->do(qq{ CREATE SCHEMA $opt{sample_schema} });
my $created_schema = 1; # keep track that we actually did it; see END block

# parse limit rules
my @limits;
foreach my $limit (grep { /\S/ } map { split /\s*,\s*/ } list($opt{limit})) {
  $limit =~ s/^\s+|\s+$//g;
  $limit = ".* = $limit" if $limit =~ /^\d+$/;

  my ($match, $rule) = ($limit =~ /^([^=\s]+) \s* = \s* (.*)$/x)
    or croak "unexpected limit value '$limit'";
  $match = '.*' if $match eq '*';

  push @limits, [$match, $rule];
}
notice "[limit] $_->[0] = $_->[1]\n" foreach @limits;

# create copies of each table in a separate schema and insert no
# more than --limit rows initially.
my @tables;
my %sample_tables; # real table name -> sample table name
my $sth = $dbh->table_info(undef, undef, undef, 'TABLE');
my $table_info = $sth->fetchall_arrayref({});
foreach my $row (@{$table_info}) {
  $row = lower_keys($row);
  next unless uc $row->{table_type} eq 'TABLE'; # skip SYSTEM TABLE values
  next if $row->{table_schem} eq 'information_schema'; # special pg schema
  next if $opt{schema} && $row->{table_schem} ne $opt{schema};

  my $sname = $row->{pg_schema} || unquote_identifier($row->{TABLE_SCHEM})
    or die "no pg_schema or TABLE_SCHEM value?!";

  my $tname = $row->{pg_table} || unquote_identifier($row->{TABLE_NAME})
    or die "no pg_table or TABLE_NAME value?!";

  my $columns = [map { $_->{column_name} } $dbh->selectall_array(qq{
      SELECT column_name
      FROM information_schema.columns
      WHERE table_name=?
      AND table_schema=?
      AND is_generated='NEVER'
      ORDER BY ordinal_position
  }, { Slice => {} }, ($tname, $sname) ) ];

  my $table = Table->new($sname, $tname, $columns);
  push @tables, $table;

  my $sample_table = sample_table($table);
  $sample_tables{ $table } = $sample_table;

  notice "Creating table $sample_table ";

  # find first matching limit rule
  my $tablesample = '';
  my $where = 'TRUE';
  my $limit = '';
  my $order = '';
  foreach (@limits) {
    $table->unquoted =~ /^$_->[0]$/i || $table->table =~ /^$_->[0]$/i or next;

    if ($_->[1] eq '*') { # include all rows
      $limit = '';
    } elsif ($_->[1] =~ /^\d+$/) { # numeric value turned into LIMIT
      if (!$opt{random} || $pg_version < version->declare('9.5')) {
        $limit = "LIMIT $_->[1]";
      } else {
        my ($table_num_rows) = $dbh->selectrow_array(qq{
          SELECT greatest(count(*), ?) FROM $table
        }, undef, $_->[1]);
        my $percent = 100 * $_->[1] / $table_num_rows;
        $tablesample = "TABLESAMPLE BERNOULLI ($percent)";
      }
    } elsif ($_->[1] =~ /^\d+(\.\d+)?%$/) { # percent value turned into LIMIT
      if (not $opt{random} or $pg_version < version->declare('9.5')) {
        my ($table_num_rows) = $dbh->selectrow_array(qq{ SELECT count(*) FROM $table });
        my $percent = (substr $_->[1], 0, (length $_->[1]) - 1) / 100;
        my $total_rows = int($table_num_rows * $percent);

        $limit = "LIMIT $total_rows";
      } else {
        my $percent = (substr $_->[1], 0, (length $_->[1]) - 1);
        $tablesample = "TABLESAMPLE BERNOULLI ($percent)";
      }
    } else { # otherwise treated as subselect
      $where = "($_->[1])";
    }

    last;
  }
  # warn "\n[LIMIT] $table WHERE $where $limit\n";

  if ($opt{random} && $pg_version < version->declare('9.5')) {
    $order = $opt{random} ? 'ORDER BY random()' : '';
  } elsif (my $direction = $opt{ordered}) {
    my @cols = find_candidate_key($table);
    if (@cols) {
      my $cols = join ', ',
        map { "$_ $direction" }
        map { $dbh->quote_identifier($_) }
        @cols
      ;
      $order = "ORDER BY $cols";
    } else {
      notice "No candidate key found for '$table'; ignoring --ordered";
    }
  }
  # Only extract non-generated columns, ORDER BY necessary because later
  # joins rely on column order to be identical between tables.
  my $quoted_cols = $table->columns_quoted;
  notice "copying cols [$quoted_cols] ";
  # Quote the column names just in case.

  $dbh->do(qq{
    CREATE $unlogged TABLE $sample_table AS
      SELECT $quoted_cols
      FROM ONLY $table
      $tablesample
      WHERE $where
      $order
      $limit
  });

  if ($opt{verbose}) {
    my ($num_rows) =
      $dbh->selectrow_array(qq{ SELECT count(*) FROM $sample_table });
    notice "$num_rows\n";
  }
}

# Find foreign keys
my @fks;
foreach my $table (@tables) {
  # find all foreign keys referencing this table
  my $sth = $dbh->foreign_key_info(
    undef, $table->schema, $table->table, # target table
    undef, undef, undef,                  # fk table
  ) or next;

  # There can be multiple rows for a single FK (composite FK)
  my %fks;
  while (my $row = lower_keys($sth->fetchrow_hashref)) {
    my $target_col = $row->{uk_column_name} or croak "no uk_column_name?!";
    my $fk_table = Table->new(
      unquote_identifier($row->{fk_table_schem}),
      unquote_identifier($row->{fk_table_name}),
      $table->columns,
    );
    my $fk_name = "$fk_table.$row->{fk_name}"; # unique key to group FK rows
    my $fk_col = $row->{fk_column_name};

    $fks{ $fk_name } ||= [$fk_table, $table];
    push @{ $fks{ $fk_name } }, [$fk_col, $target_col];
  }
  push @fks, values %fks;
}

# create indexes on sample table FKs
my $idx = 0;
foreach my $fk (@fks) {
  my ($fk_table, $table, @pairs) = @$fk;

  my $sample_fk_table = $sample_tables{ $fk_table };
  my $idx_name = $dbh->quote_identifier("$opt{sample_schema}_idx" . ++$idx);
  my $fk_cols = join ', ', map { $_->[0] } @pairs;
  $dbh->do(qq{ CREATE INDEX $idx_name ON $sample_fk_table ($fk_cols) });
}


# Skip FK inserts if entire table is being included anyway
# (useful for cases where we don't support equality operator)
sub should_skip_insert_for_fk ($) {
  my $table_name = shift;
  
  $table_name =~ s/"\."/\./g;
  notice "$table_name\n";
  foreach my $limit (@limits) {
    my ($regex, $action) = @$limit;
    if ($table_name =~ $regex && $action eq '*') {
      notice "Skipping $table_name as all data is imported\n";
      return 1;  # Skip if it matches the regex and action is '*'
    }
  }
  return 0;  # Do not skip by default
}


# Keep inserting rows to satisfy any fk constraints until no more
# are inserted. This should handle circular references.
my $num_rows = 1;
while ($num_rows) {
  $num_rows = 0;

  foreach my $fk (@fks) {
    my ($fk_table, $target_table, @pairs) = @$fk;

    next if should_skip_insert_for_fk($target_table);

    my $target_sample_table = $sample_tables{ $target_table };
    my $sample_fk_table = $sample_tables{ $fk_table };

    my $fk_cols = join ', ', map { $_->[0] } @pairs;
    my $target_cols = join ', ', map { $_->[1] } @pairs;

    my $join1 = join ' AND ',
      map { "f1.$_->[0] = t1.$_->[1]" } @pairs;

    my $join2 = join ' AND ',
      map { "t1.$_->[1] = s1.$_->[1]" } @pairs;

    my $where = join ' AND ',
      map { "s1.$_->[1] IS NULL" } @pairs;

    # Insert into the sample table all the rows needed to
    # satisfy the fk table, except those already present.
    my $quoted_cols = $target_table->columns_quoted("t1");
    my $query = qq{
      INSERT INTO $target_sample_table
           SELECT DISTINCT $quoted_cols
             FROM $target_table t1
                  JOIN $sample_fk_table f1 ON ($join1)
                  LEFT JOIN $target_sample_table s1 ON ($join2)
            WHERE $where
    };

    #warn "\nQUERY:\n$query\n";
    # my $sth = $dbh->prepare(qq{ EXPLAIN $query });
    # $sth->execute;
    # while (my $row = $sth->fetchrow_arrayref) {
    #   warn @$row, "\n";
    # }
    # warn "\n\n";

    notice
      "Copying $target_table ($target_cols) rows referenced from " .
      "$sample_fk_table ($fk_cols)... "
    ;

    my $count = $dbh->do($query) || 0;
    notice "$count rows\n";

    $num_rows += $count;
  }
}

# fetch all sequences and current values
$sth = $dbh->prepare(qq{
  SELECT sequence_schema
         ,sequence_name
    FROM information_schema.sequences
   WHERE sequence_schema NOT LIKE 'pg_%'
});
$sth->execute;
my %seq;
while (my $row = $sth->fetchrow_hashref) {
  next if $opt{schema} && $row->{sequence_schema} ne $opt{schema};

  my $name = Table->new($row->{sequence_schema}, $row->{sequence_name});
  $seq{ $name } = 0;
}
foreach my $name (keys %seq) {
  my ($lastval) = $dbh->selectrow_array(qq{ SELECT last_value FROM $name });
  $seq{ $name } = $lastval;
}

# Most of these were copied from the pg_dump output.
print <<EOF;

SET client_encoding = '$opt{encoding}';
SET standard_conforming_strings = on;
SET check_function_bodies = false;
SET client_min_messages = warning;
SET escape_string_warning = off;

EOF

notice "Exporting sequences\n";
print "\n";
foreach my $name (sort keys %seq) {
  my $constant = quote_constant($name);
  print "SELECT pg_catalog.setval($constant, $seq{ $name });\n";
}
print "\n";

# Disable all foreign key constraints temporarily
print "\n";
foreach my $table (@tables) {
  print "ALTER TABLE $table DISABLE TRIGGER ALL;\n";
}
print "\n";

foreach my $table (@tables) {
  my $sample_table = sample_table($table);
  if ($opt{verbose}) {
    my ($count) = $dbh->selectrow_array("SELECT count(*) FROM $sample_table");
    notice "Exporting data from $sample_table ($count)\n";
  }
  # Specify column ordering explicitly in COPY FROM in case the ordering is
  # different in the restored database.
  my $columns = $sample_table->columns_quoted;
  print "COPY $table($columns) FROM stdin;\n";

  if ($opt{ordered}) {
    my @cols = find_candidate_key($table);
    if (@cols) {
      my $cols = join ', ', map { $dbh->quote_identifier($_) } @cols;
      $sample_table = "(SELECT * FROM $sample_table ORDER BY $cols)";
    } else {
      notice "No candidate key found for '$table'; ignoring --ordered";
    }
  }

  $dbh->do(qq{ COPY $sample_table TO STDOUT });
  my $buffer = '';
  print $buffer while $dbh->pg_getcopydata($buffer) >= 0;
  print "\\.\n\n";
}

# Re-enable all triggers
print "\n";
foreach my $table (@tables) {
  print "ALTER TABLE $table ENABLE TRIGGER ALL;\n";
}
print "\n";

END {
  # remove sample tables unless requested not to
  if ($created_schema && !$opt{keep}) {
    notice "Dropping sample schema $opt{sample_schema}\n";
    $dbh->do("DROP SCHEMA $opt{sample_schema} CASCADE");
  }
  notice "Done.\n";

  exit $exit_code;
}