Multi: Add pg_cmd function for more robust error handling

2 files changed, 49 insertions, 20 deletions

diff --git a/lib/Multi/Core.pm b/lib/Multi/Core.pm
index 5bfa2a19..ac58403c 100644
--- a/lib/Multi/Core.pm
+++ b/lib/Multi/Core.pm
@@ -15,7 +15,7 @@ use DBI;
 use POSIX 'setsid', 'pause', 'SIGUSR1';
 use Exporter 'import';
 
-our @EXPORT = qw|pg pg_expect schedule push_watcher|;
+our @EXPORT = qw|pg pg_cmd pg_expect schedule push_watcher|;
 
 
 my $PG;
@@ -76,6 +76,7 @@ sub load_pg {
   my %vars = split /[,=]/, $dsn[4];
   $PG = AnyEvent::Pg::Pool->new(
     {%vars, user => $db[1], password => $db[2], host => 'localhost'},
+    timeout => 0, # Some maintenance queries can take a while to run...
     on_error => sub { die "Lost connection to PostgreSQL\n"; },
     on_connect_error => sub { die "Lost connection to PostgreSQL\n"; },
   );
@@ -151,15 +152,49 @@ sub schedule {
 # Logs any unexpected results and returns 0 if the expectations were met.
 sub pg_expect {
   my($res, $exp) = @_;
-  return 0 if !$exp && $res->status == PGRES_COMMAND_OK;
-  return 0 if $exp && $res->status == PGRES_TUPLES_OK;
-  AE::log alert => $res->errorMessage
+  return 0 if !$exp && $res && $res->status == PGRES_COMMAND_OK;
+  return 0 if $exp && $res && $res->status == PGRES_TUPLES_OK;
+  AE::log alert => !$res
+    ? sprintf 'AnyEvent::Pg error at %s:%d', (caller)[0,2] : $res->errorMessage
     ? sprintf 'SQL error at %s:%d: %s', (caller)[0,2], $res->errorMessage
     : sprintf 'Unexpected status at %s:%d: %s', (caller)[0,2], $res->statusMessage;
   return 1;
 }
 
 
+# Wrapper around pg->push_query().
+# Args: $query, \@args, sub {}
+# The sub will be called on either on_error or on_done, and has two args: The
+# result and the running time. Only a single on_result is expected. The result
+# argument is undef on error.
+# Unlike most AE watchers, this function does not return a watcher object and
+# can not be cancelled.
+sub pg_cmd {
+  my($q, $a, $s) = @_;
+  my $r;
+  my $w; $w = pg->push_query(
+    query => $q,
+    $a ? (args => $a) : (),
+    on_error => sub {
+      undef $w;
+      $s->(undef, 0);
+    },
+    on_result => sub {
+      if($r) {
+        AE::log warn => "Received more than one result for query: $q";
+        $s->(undef, 0);
+      } else {
+        $r = $_[2];
+      }
+    },
+    on_done => sub {
+      undef $w;
+      $s->($r, $_[2]);
+    },
+  );
+}
+
+
 # Tiny class for forwarding output for STDERR/STDOUT to the log file using tie().
 package Multi::Core::STDIO;
 
diff --git a/lib/Multi/Feed.pm b/lib/Multi/Feed.pm
index ac171c15..35807e67 100644
--- a/lib/Multi/Feed.pm
+++ b/lib/Multi/Feed.pm
@@ -29,8 +29,7 @@ sub run {
 
 sub generate {
   # announcements
-  my $a; $a = pg->push_query(
-    query => q{
+  pg_cmd q{
       SELECT '/t'||t.id AS id, t.title, extract('epoch' from tp.date) AS published,
          extract('epoch' from tp.edited) AS updated, u.username, u.id AS uid, tp.msg AS summary
        FROM threads t
@@ -40,13 +39,11 @@ sub generate {
       WHERE NOT t.hidden
       ORDER BY t.id DESC
       LIMIT $1},
-    args => [$VNDB::S{atom_feeds}{announcements}[0]],
-    on_result => sub { write_atom(announcements => @_[2,3], $a) },
-  );
+    [$VNDB::S{atom_feeds}{announcements}[0]],
+    sub { write_atom(announcements => @_) };
 
   # changes
-  my $c; $c = pg->push_query(
-    query => q{
+  pg_cmd q{
       SELECT '/'||c.type||COALESCE(vr.vid, rr.rid, pr.pid, cr.cid)||'.'||c.rev AS id,
          COALESCE(vr.title, rr.title, pr.name, cr.name) AS title, extract('epoch' from c.added) AS updated,
          u.username, u.id AS uid, c.comments AS summary
@@ -59,13 +56,11 @@ sub generate {
       WHERE c.requester <> 1
       ORDER BY c.id DESC
       LIMIT $1},
-    args => [$VNDB::S{atom_feeds}{changes}[0]],
-    on_result => sub { write_atom(changes => @_[2,3], $c); },
-  );
+    [$VNDB::S{atom_feeds}{changes}[0]],
+    sub { write_atom(changes => @_); };
 
   # posts (this query isn't all that fast)
-  my $p; $p = pg->push_query(
-    query => q{
+  pg_cmd q{
       SELECT '/t'||t.id||'.'||tp.num AS id, t.title||' (#'||tp.num||')' AS title, extract('epoch' from tp.date) AS published,
          extract('epoch' from tp.edited) AS updated, u.username, u.id AS uid, tp.msg AS summary
        FROM threads_posts tp
@@ -73,10 +68,9 @@ sub generate {
        JOIN users u ON u.id = tp.uid
       WHERE NOT tp.hidden AND NOT t.hidden
       ORDER BY tp.date DESC
-      LIMIT $1},
-    args => [$VNDB::S{atom_feeds}{posts}[0]],
-    on_result => sub { write_atom(posts => @_[2,3], $p); },
-  );
+      LIMIT ?$1},
+    [$VNDB::S{atom_feeds}{posts}[0]],
+    sub { write_atom(posts => @_); };
 }