granicus.if.org Git - postgresql/blob - src/bin/pg_dump/pg_backup_archiver.c

   1 /*-------------------------------------------------------------------------
   2  *
   3  * pg_backup_archiver.c
   4  *
   5  *      Private implementation of the archiver routines.
   6  *
   7  *      See the headers to pg_restore for more details.
   8  *
   9  * Copyright (c) 2000, Philip Warner
  10  *      Rights are granted to use this software in any way so long
  11  *      as this notice is not removed.
  12  *
  13  *      The author is not responsible for loss or damages that may
  14  *      result from its use.
  15  *
  16  *
  17  * IDENTIFICATION
  18  *              $PostgreSQL: pgsql/src/bin/pg_dump/pg_backup_archiver.c,v 1.187 2010/07/06 19:18:59 momjian Exp $
  19  *
  20  *-------------------------------------------------------------------------
  21  */
  22
  23 #include "pg_backup_db.h"
  24 #include "dumputils.h"
  25
  26 #include <ctype.h>
  27 #include <unistd.h>
  28 #include <sys/types.h>
  29 #include <sys/wait.h>
  30
  31 #ifdef WIN32
  32 #include <io.h>
  33 #endif
  34
  35 #include "libpq/libpq-fs.h"
  36
  37 /*
  38  * Special exit values from worker children.  We reserve 0 for normal
  39  * success; 1 and other small values should be interpreted as crashes.
  40  */
  41 #define WORKER_CREATE_DONE              10
  42 #define WORKER_INHIBIT_DATA             11
  43 #define WORKER_IGNORED_ERRORS   12
  44
  45 /*
  46  * Unix uses exit to return result from worker child, so function is void.
  47  * Windows thread result comes via function return.
  48  */
  49 #ifndef WIN32
  50 #define parallel_restore_result void
  51 #else
  52 #define parallel_restore_result DWORD
  53 #endif
  54
  55 /* IDs for worker children are either PIDs or thread handles */
  56 #ifndef WIN32
  57 #define thandle pid_t
  58 #else
  59 #define thandle HANDLE
  60 #endif
  61
  62 /* Arguments needed for a worker child */
  63 typedef struct _restore_args
  64 {
  65         ArchiveHandle *AH;
  66         TocEntry   *te;
  67 } RestoreArgs;
  68
  69 /* State for each parallel activity slot */
  70 typedef struct _parallel_slot
  71 {
  72         thandle         child_id;
  73         RestoreArgs *args;
  74 } ParallelSlot;
  75
  76 #define NO_SLOT (-1)
  77
  78 const char *progname;
  79
  80 static const char *modulename = gettext_noop("archiver");
  81
  82
  83 static ArchiveHandle *_allocAH(const char *FileSpec, const ArchiveFormat fmt,
  84                  const int compression, ArchiveMode mode);
  85 static void _getObjectDescription(PQExpBuffer buf, TocEntry *te,
  86                                           ArchiveHandle *AH);
  87 static void _printTocEntry(ArchiveHandle *AH, TocEntry *te, RestoreOptions *ropt, bool isData, bool acl_pass);
  88
  89
  90 static void _doSetFixedOutputState(ArchiveHandle *AH);
  91 static void _doSetSessionAuth(ArchiveHandle *AH, const char *user);
  92 static void _doSetWithOids(ArchiveHandle *AH, const bool withOids);
  93 static void _reconnectToDB(ArchiveHandle *AH, const char *dbname);
  94 static void _becomeUser(ArchiveHandle *AH, const char *user);
  95 static void _becomeOwner(ArchiveHandle *AH, TocEntry *te);
  96 static void _selectOutputSchema(ArchiveHandle *AH, const char *schemaName);
  97 static void _selectTablespace(ArchiveHandle *AH, const char *tablespace);
  98 static void processEncodingEntry(ArchiveHandle *AH, TocEntry *te);
  99 static void processStdStringsEntry(ArchiveHandle *AH, TocEntry *te);
 100 static teReqs _tocEntryRequired(TocEntry *te, RestoreOptions *ropt, bool include_acls);
 101 static bool _tocEntryIsACL(TocEntry *te);
 102 static void _disableTriggersIfNecessary(ArchiveHandle *AH, TocEntry *te, RestoreOptions *ropt);
 103 static void _enableTriggersIfNecessary(ArchiveHandle *AH, TocEntry *te, RestoreOptions *ropt);
 104 static TocEntry *getTocEntryByDumpId(ArchiveHandle *AH, DumpId id);
 105 static void _moveAfter(ArchiveHandle *AH, TocEntry *pos, TocEntry *te);
 106 static int      _discoverArchiveFormat(ArchiveHandle *AH);
 107
 108 static void dump_lo_buf(ArchiveHandle *AH);
 109 static void _write_msg(const char *modulename, const char *fmt, va_list ap);
 110 static void _die_horribly(ArchiveHandle *AH, const char *modulename, const char *fmt, va_list ap);
 111
 112 static void dumpTimestamp(ArchiveHandle *AH, const char *msg, time_t tim);
 113 static OutputContext SetOutput(ArchiveHandle *AH, char *filename, int compression);
 114 static void ResetOutput(ArchiveHandle *AH, OutputContext savedContext);
 115
 116 static int restore_toc_entry(ArchiveHandle *AH, TocEntry *te,
 117                                   RestoreOptions *ropt, bool is_parallel);
 118 static void restore_toc_entries_parallel(ArchiveHandle *AH);
 119 static thandle spawn_restore(RestoreArgs *args);
 120 static thandle reap_child(ParallelSlot *slots, int n_slots, int *work_status);
 121 static bool work_in_progress(ParallelSlot *slots, int n_slots);
 122 static int      get_next_slot(ParallelSlot *slots, int n_slots);
 123 static void par_list_header_init(TocEntry *l);
 124 static void par_list_append(TocEntry *l, TocEntry *te);
 125 static void par_list_remove(TocEntry *te);
 126 static TocEntry *get_next_work_item(ArchiveHandle *AH,
 127                                    TocEntry *ready_list,
 128                                    ParallelSlot *slots, int n_slots);
 129 static parallel_restore_result parallel_restore(RestoreArgs *args);
 130 static void mark_work_done(ArchiveHandle *AH, TocEntry *ready_list,
 131                            thandle worker, int status,
 132                            ParallelSlot *slots, int n_slots);
 133 static void fix_dependencies(ArchiveHandle *AH);
 134 static bool has_lock_conflicts(TocEntry *te1, TocEntry *te2);
 135 static void repoint_table_dependencies(ArchiveHandle *AH,
 136                                                    DumpId tableId, DumpId tableDataId);
 137 static void identify_locking_dependencies(TocEntry *te,
 138                                                           TocEntry **tocsByDumpId,
 139                                                           DumpId maxDumpId);
 140 static void reduce_dependencies(ArchiveHandle *AH, TocEntry *te,
 141                                         TocEntry *ready_list);
 142 static void mark_create_done(ArchiveHandle *AH, TocEntry *te);
 143 static void inhibit_data_for_failed_table(ArchiveHandle *AH, TocEntry *te);
 144 static ArchiveHandle *CloneArchive(ArchiveHandle *AH);
 145 static void DeCloneArchive(ArchiveHandle *AH);
 146
 147
 148 /*
 149  *      Wrapper functions.
 150  *
 151  *      The objective it to make writing new formats and dumpers as simple
 152  *      as possible, if necessary at the expense of extra function calls etc.
 153  *
 154  */
 155
 156
 157 /* Create a new archive */
 158 /* Public */
 159 Archive *
 160 CreateArchive(const char *FileSpec, const ArchiveFormat fmt,
 161                           const int compression, ArchiveMode mode)
 162
 163 {
 164         ArchiveHandle *AH = _allocAH(FileSpec, fmt, compression, mode);
 165
 166         return (Archive *) AH;
 167 }
 168
 169 /* Open an existing archive */
 170 /* Public */
 171 Archive *
 172 OpenArchive(const char *FileSpec, const ArchiveFormat fmt)
 173 {
 174         ArchiveHandle *AH = _allocAH(FileSpec, fmt, 0, archModeRead);
 175
 176         return (Archive *) AH;
 177 }
 178
 179 /* Public */
 180 void
 181 CloseArchive(Archive *AHX)
 182 {
 183         int                     res = 0;
 184         ArchiveHandle *AH = (ArchiveHandle *) AHX;
 185
 186         (*AH->ClosePtr) (AH);
 187
 188         /* Close the output */
 189         if (AH->gzOut)
 190                 res = GZCLOSE(AH->OF);
 191         else if (AH->OF != stdout)
 192                 res = fclose(AH->OF);
 193
 194         if (res != 0)
 195                 die_horribly(AH, modulename, "could not close output file: %s\n",
 196                                          strerror(errno));
 197 }
 198
 199 /* Public */
 200 void
 201 RestoreArchive(Archive *AHX, RestoreOptions *ropt)
 202 {
 203         ArchiveHandle *AH = (ArchiveHandle *) AHX;
 204         TocEntry   *te;
 205         teReqs          reqs;
 206         OutputContext sav;
 207
 208         AH->ropt = ropt;
 209         AH->stage = STAGE_INITIALIZING;
 210
 211         /*
 212          * Check for nonsensical option combinations.
 213          *
 214          * NB: createDB+dropSchema is useless because if you're creating the DB,
 215          * there's no need to drop individual items in it.  Moreover, if we tried
 216          * to do that then we'd issue the drops in the database initially
 217          * connected to, not the one we will create, which is very bad...
 218          */
 219         if (ropt->createDB && ropt->dropSchema)
 220                 die_horribly(AH, modulename, "-C and -c are incompatible options\n");
 221
 222         /*
 223          * -C is not compatible with -1, because we can't create a database inside
 224          * a transaction block.
 225          */
 226         if (ropt->createDB && ropt->single_txn)
 227                 die_horribly(AH, modulename, "-C and -1 are incompatible options\n");
 228
 229         /*
 230          * Make sure we won't need (de)compression we haven't got
 231          */
 232 #ifndef HAVE_LIBZ
 233         if (AH->compression != 0 && AH->PrintTocDataPtr !=NULL)
 234         {
 235                 for (te = AH->toc->next; te != AH->toc; te = te->next)
 236                 {
 237                         reqs = _tocEntryRequired(te, ropt, false);
 238                         if (te->hadDumper && (reqs & REQ_DATA) != 0)
 239                                 die_horribly(AH, modulename, "cannot restore from compressed archive (compression not supported in this installation)\n");
 240                 }
 241         }
 242 #endif
 243
 244         /*
 245          * If we're using a DB connection, then connect it.
 246          */
 247         if (ropt->useDB)
 248         {
 249                 ahlog(AH, 1, "connecting to database for restore\n");
 250                 if (AH->version < K_VERS_1_3)
 251                         die_horribly(AH, modulename, "direct database connections are not supported in pre-1.3 archives\n");
 252
 253                 /* XXX Should get this from the archive */
 254                 AHX->minRemoteVersion = 070100;
 255                 AHX->maxRemoteVersion = 999999;
 256
 257                 ConnectDatabase(AHX, ropt->dbname,
 258                                                 ropt->pghost, ropt->pgport, ropt->username,
 259                                                 ropt->promptPassword);
 260
 261                 /*
 262                  * If we're talking to the DB directly, don't send comments since they
 263                  * obscure SQL when displaying errors
 264                  */
 265                 AH->noTocComments = 1;
 266         }
 267
 268         /*
 269          * Work out if we have an implied data-only restore. This can happen if
 270          * the dump was data only or if the user has used a toc list to exclude
 271          * all of the schema data. All we do is look for schema entries - if none
 272          * are found then we set the dataOnly flag.
 273          *
 274          * We could scan for wanted TABLE entries, but that is not the same as
 275          * dataOnly. At this stage, it seems unnecessary (6-Mar-2001).
 276          */
 277         if (!ropt->dataOnly)
 278         {
 279                 int                     impliedDataOnly = 1;
 280
 281                 for (te = AH->toc->next; te != AH->toc; te = te->next)
 282                 {
 283                         reqs = _tocEntryRequired(te, ropt, true);
 284                         if ((reqs & REQ_SCHEMA) != 0)
 285                         {                                       /* It's schema, and it's wanted */
 286                                 impliedDataOnly = 0;
 287                                 break;
 288                         }
 289                 }
 290                 if (impliedDataOnly)
 291                 {
 292                         ropt->dataOnly = impliedDataOnly;
 293                         ahlog(AH, 1, "implied data-only restore\n");
 294                 }
 295         }
 296
 297         /*
 298          * Setup the output file if necessary.
 299          */
 300         if (ropt->filename || ropt->compression)
 301                 sav = SetOutput(AH, ropt->filename, ropt->compression);
 302
 303         ahprintf(AH, "--\n-- PostgreSQL database dump\n--\n\n");
 304
 305         if (AH->public.verbose)
 306         {
 307                 if (AH->archiveRemoteVersion)
 308                         ahprintf(AH, "-- Dumped from database version %s\n",
 309                                          AH->archiveRemoteVersion);
 310                 if (AH->archiveDumpVersion)
 311                         ahprintf(AH, "-- Dumped by pg_dump version %s\n",
 312                                          AH->archiveDumpVersion);
 313                 dumpTimestamp(AH, "Started on", AH->createDate);
 314         }
 315
 316         if (ropt->single_txn)
 317         {
 318                 if (AH->connection)
 319                         StartTransaction(AH);
 320                 else
 321                         ahprintf(AH, "BEGIN;\n\n");
 322         }
 323
 324         /*
 325          * Establish important parameter values right away.
 326          */
 327         _doSetFixedOutputState(AH);
 328
 329         AH->stage = STAGE_PROCESSING;
 330
 331         /*
 332          * Drop the items at the start, in reverse order
 333          */
 334         if (ropt->dropSchema)
 335         {
 336                 for (te = AH->toc->prev; te != AH->toc; te = te->prev)
 337                 {
 338                         AH->currentTE = te;
 339
 340                         reqs = _tocEntryRequired(te, ropt, false /* needn't drop ACLs */ );
 341                         /* We want anything that's selected and has a dropStmt */
 342                         if (((reqs & (REQ_SCHEMA | REQ_DATA)) != 0) && te->dropStmt)
 343                         {
 344                                 ahlog(AH, 1, "dropping %s %s\n", te->desc, te->tag);
 345                                 /* Select owner and schema as necessary */
 346                                 _becomeOwner(AH, te);
 347                                 _selectOutputSchema(AH, te->namespace);
 348                                 /* Drop it */
 349                                 ahprintf(AH, "%s", te->dropStmt);
 350                         }
 351                 }
 352
 353                 /*
 354                  * _selectOutputSchema may have set currSchema to reflect the effect
 355                  * of a "SET search_path" command it emitted.  However, by now we may
 356                  * have dropped that schema; or it might not have existed in the first
 357                  * place.  In either case the effective value of search_path will not
 358                  * be what we think.  Forcibly reset currSchema so that we will
 359                  * re-establish the search_path setting when needed (after creating
 360                  * the schema).
 361                  *
 362                  * If we treated users as pg_dump'able objects then we'd need to reset
 363                  * currUser here too.
 364                  */
 365                 if (AH->currSchema)
 366                         free(AH->currSchema);
 367                 AH->currSchema = NULL;
 368         }
 369
 370         /*
 371          * In serial mode, we now process each non-ACL TOC entry.
 372          *
 373          * In parallel mode, turn control over to the parallel-restore logic.
 374          */
 375         if (ropt->number_of_jobs > 1 && ropt->useDB)
 376                 restore_toc_entries_parallel(AH);
 377         else
 378         {
 379                 for (te = AH->toc->next; te != AH->toc; te = te->next)
 380                         (void) restore_toc_entry(AH, te, ropt, false);
 381         }
 382
 383         /*
 384          * Scan TOC again to output ownership commands and ACLs
 385          */
 386         for (te = AH->toc->next; te != AH->toc; te = te->next)
 387         {
 388                 AH->currentTE = te;
 389
 390                 /* Work out what, if anything, we want from this entry */
 391                 reqs = _tocEntryRequired(te, ropt, true);
 392
 393                 /* Both schema and data objects might now have ownership/ACLs */
 394                 if ((reqs & (REQ_SCHEMA | REQ_DATA)) != 0)
 395                 {
 396                         ahlog(AH, 1, "setting owner and privileges for %s %s\n",
 397                                   te->desc, te->tag);
 398                         _printTocEntry(AH, te, ropt, false, true);
 399                 }
 400         }
 401
 402         if (ropt->single_txn)
 403         {
 404                 if (AH->connection)
 405                         CommitTransaction(AH);
 406                 else
 407                         ahprintf(AH, "COMMIT;\n\n");
 408         }
 409
 410         if (AH->public.verbose)
 411                 dumpTimestamp(AH, "Completed on", time(NULL));
 412
 413         ahprintf(AH, "--\n-- PostgreSQL database dump complete\n--\n\n");
 414
 415         /*
 416          * Clean up & we're done.
 417          */
 418         AH->stage = STAGE_FINALIZING;
 419
 420         if (ropt->filename || ropt->compression)
 421                 ResetOutput(AH, sav);
 422
 423         if (ropt->useDB)
 424         {
 425                 PQfinish(AH->connection);
 426                 AH->connection = NULL;
 427         }
 428 }
 429
 430 /*
 431  * Restore a single TOC item.  Used in both parallel and non-parallel restore;
 432  * is_parallel is true if we are in a worker child process.
 433  *
 434  * Returns 0 normally, but WORKER_CREATE_DONE or WORKER_INHIBIT_DATA if
 435  * the parallel parent has to make the corresponding status update.
 436  */
 437 static int
 438 restore_toc_entry(ArchiveHandle *AH, TocEntry *te,
 439                                   RestoreOptions *ropt, bool is_parallel)
 440 {
 441         int                     retval = 0;
 442         teReqs          reqs;
 443         bool            defnDumped;
 444
 445         AH->currentTE = te;
 446
 447         /* Work out what, if anything, we want from this entry */
 448         reqs = _tocEntryRequired(te, ropt, false);
 449
 450         /* Dump any relevant dump warnings to stderr */
 451         if (!ropt->suppressDumpWarnings && strcmp(te->desc, "WARNING") == 0)
 452         {
 453                 if (!ropt->dataOnly && te->defn != NULL && strlen(te->defn) != 0)
 454                         write_msg(modulename, "warning from original dump file: %s\n", te->defn);
 455                 else if (te->copyStmt != NULL && strlen(te->copyStmt) != 0)
 456                         write_msg(modulename, "warning from original dump file: %s\n", te->copyStmt);
 457         }
 458
 459         defnDumped = false;
 460
 461         if ((reqs & REQ_SCHEMA) != 0)           /* We want the schema */
 462         {
 463                 ahlog(AH, 1, "creating %s %s\n", te->desc, te->tag);
 464
 465                 _printTocEntry(AH, te, ropt, false, false);
 466                 defnDumped = true;
 467
 468                 if (strcmp(te->desc, "TABLE") == 0)
 469                 {
 470                         if (AH->lastErrorTE == te)
 471                         {
 472                                 /*
 473                                  * We failed to create the table. If
 474                                  * --no-data-for-failed-tables was given, mark the
 475                                  * corresponding TABLE DATA to be ignored.
 476                                  *
 477                                  * In the parallel case this must be done in the parent, so we
 478                                  * just set the return value.
 479                                  */
 480                                 if (ropt->noDataForFailedTables)
 481                                 {
 482                                         if (is_parallel)
 483                                                 retval = WORKER_INHIBIT_DATA;
 484                                         else
 485                                                 inhibit_data_for_failed_table(AH, te);
 486                                 }
 487                         }
 488                         else
 489                         {
 490                                 /*
 491                                  * We created the table successfully.  Mark the corresponding
 492                                  * TABLE DATA for possible truncation.
 493                                  *
 494                                  * In the parallel case this must be done in the parent, so we
 495                                  * just set the return value.
 496                                  */
 497                                 if (is_parallel)
 498                                         retval = WORKER_CREATE_DONE;
 499                                 else
 500                                         mark_create_done(AH, te);
 501                         }
 502                 }
 503
 504                 /* If we created a DB, connect to it... */
 505                 if (strcmp(te->desc, "DATABASE") == 0)
 506                 {
 507                         ahlog(AH, 1, "connecting to new database \"%s\"\n", te->tag);
 508                         _reconnectToDB(AH, te->tag);
 509                         ropt->dbname = strdup(te->tag);
 510                 }
 511         }
 512
 513         /*
 514          * If we have a data component, then process it
 515          */
 516         if ((reqs & REQ_DATA) != 0)
 517         {
 518                 /*
 519                  * hadDumper will be set if there is genuine data component for this
 520                  * node. Otherwise, we need to check the defn field for statements
 521                  * that need to be executed in data-only restores.
 522                  */
 523                 if (te->hadDumper)
 524                 {
 525                         /*
 526                          * If we can output the data, then restore it.
 527                          */
 528                         if (AH->PrintTocDataPtr !=NULL && (reqs & REQ_DATA) != 0)
 529                         {
 530                                 _printTocEntry(AH, te, ropt, true, false);
 531
 532                                 if (strcmp(te->desc, "BLOBS") == 0 ||
 533                                         strcmp(te->desc, "BLOB COMMENTS") == 0)
 534                                 {
 535                                         ahlog(AH, 1, "restoring %s\n", te->desc);
 536
 537                                         _selectOutputSchema(AH, "pg_catalog");
 538
 539                                         (*AH->PrintTocDataPtr) (AH, te, ropt);
 540                                 }
 541                                 else
 542                                 {
 543                                         _disableTriggersIfNecessary(AH, te, ropt);
 544
 545                                         /* Select owner and schema as necessary */
 546                                         _becomeOwner(AH, te);
 547                                         _selectOutputSchema(AH, te->namespace);
 548
 549                                         ahlog(AH, 1, "restoring data for table \"%s\"\n",
 550                                                   te->tag);
 551
 552                                         /*
 553                                          * In parallel restore, if we created the table earlier in
 554                                          * the run then we wrap the COPY in a transaction and
 555                                          * precede it with a TRUNCATE.  If archiving is not on
 556                                          * this prevents WAL-logging the COPY.  This obtains a
 557                                          * speedup similar to that from using single_txn mode in
 558                                          * non-parallel restores.
 559                                          */
 560                                         if (is_parallel && te->created)
 561                                         {
 562                                                 /*
 563                                                  * Parallel restore is always talking directly to a
 564                                                  * server, so no need to see if we should issue BEGIN.
 565                                                  */
 566                                                 StartTransaction(AH);
 567
 568                                                 /*
 569                                                  * If the server version is >= 8.4, make sure we issue
 570                                                  * TRUNCATE with ONLY so that child tables are not
 571                                                  * wiped.
 572                                                  */
 573                                                 ahprintf(AH, "TRUNCATE TABLE %s%s;\n\n",
 574                                                                  (PQserverVersion(AH->connection) >= 80400 ?
 575                                                                   "ONLY " : ""),
 576                                                                  fmtId(te->tag));
 577                                         }
 578
 579                                         /*
 580                                          * If we have a copy statement, use it. As of V1.3, these
 581                                          * are separate to allow easy import from withing a
 582                                          * database connection. Pre 1.3 archives can not use DB
 583                                          * connections and are sent to output only.
 584                                          *
 585                                          * For V1.3+, the table data MUST have a copy statement so
 586                                          * that we can go into appropriate mode with libpq.
 587                                          */
 588                                         if (te->copyStmt && strlen(te->copyStmt) > 0)
 589                                         {
 590                                                 ahprintf(AH, "%s", te->copyStmt);
 591                                                 AH->writingCopyData = true;
 592                                         }
 593
 594                                         (*AH->PrintTocDataPtr) (AH, te, ropt);
 595
 596                                         AH->writingCopyData = false;
 597
 598                                         /* close out the transaction started above */
 599                                         if (is_parallel && te->created)
 600                                                 CommitTransaction(AH);
 601
 602                                         _enableTriggersIfNecessary(AH, te, ropt);
 603                                 }
 604                         }
 605                 }
 606                 else if (!defnDumped)
 607                 {
 608                         /* If we haven't already dumped the defn part, do so now */
 609                         ahlog(AH, 1, "executing %s %s\n", te->desc, te->tag);
 610                         _printTocEntry(AH, te, ropt, false, false);
 611                 }
 612         }
 613
 614         return retval;
 615 }
 616
 617 /*
 618  * Allocate a new RestoreOptions block.
 619  * This is mainly so we can initialize it, but also for future expansion,
 620  */
 621 RestoreOptions *
 622 NewRestoreOptions(void)
 623 {
 624         RestoreOptions *opts;
 625
 626         opts = (RestoreOptions *) calloc(1, sizeof(RestoreOptions));
 627
 628         /* set any fields that shouldn't default to zeroes */
 629         opts->format = archUnknown;
 630         opts->promptPassword = TRI_DEFAULT;
 631
 632         return opts;
 633 }
 634
 635 static void
 636 _disableTriggersIfNecessary(ArchiveHandle *AH, TocEntry *te, RestoreOptions *ropt)
 637 {
 638         /* This hack is only needed in a data-only restore */
 639         if (!ropt->dataOnly || !ropt->disable_triggers)
 640                 return;
 641
 642         ahlog(AH, 1, "disabling triggers for %s\n", te->tag);
 643
 644         /*
 645          * Become superuser if possible, since they are the only ones who can
 646          * disable constraint triggers.  If -S was not given, assume the initial
 647          * user identity is a superuser.  (XXX would it be better to become the
 648          * table owner?)
 649          */
 650         _becomeUser(AH, ropt->superuser);
 651
 652         /*
 653          * Disable them.
 654          */
 655         _selectOutputSchema(AH, te->namespace);
 656
 657         ahprintf(AH, "ALTER TABLE %s DISABLE TRIGGER ALL;\n\n",
 658                          fmtId(te->tag));
 659 }
 660
 661 static void
 662 _enableTriggersIfNecessary(ArchiveHandle *AH, TocEntry *te, RestoreOptions *ropt)
 663 {
 664         /* This hack is only needed in a data-only restore */
 665         if (!ropt->dataOnly || !ropt->disable_triggers)
 666                 return;
 667
 668         ahlog(AH, 1, "enabling triggers for %s\n", te->tag);
 669
 670         /*
 671          * Become superuser if possible, since they are the only ones who can
 672          * disable constraint triggers.  If -S was not given, assume the initial
 673          * user identity is a superuser.  (XXX would it be better to become the
 674          * table owner?)
 675          */
 676         _becomeUser(AH, ropt->superuser);
 677
 678         /*
 679          * Enable them.
 680          */
 681         _selectOutputSchema(AH, te->namespace);
 682
 683         ahprintf(AH, "ALTER TABLE %s ENABLE TRIGGER ALL;\n\n",
 684                          fmtId(te->tag));
 685 }
 686
 687 /*
 688  * This is a routine that is part of the dumper interface, hence the 'Archive*' parameter.
 689  */
 690
 691 /* Public */
 692 size_t
 693 WriteData(Archive *AHX, const void *data, size_t dLen)
 694 {
 695         ArchiveHandle *AH = (ArchiveHandle *) AHX;
 696
 697         if (!AH->currToc)
 698                 die_horribly(AH, modulename, "internal error -- WriteData cannot be called outside the context of a DataDumper routine\n");
 699
 700         return (*AH->WriteDataPtr) (AH, data, dLen);
 701 }
 702
 703 /*
 704  * Create a new TOC entry. The TOC was designed as a TOC, but is now the
 705  * repository for all metadata. But the name has stuck.
 706  */
 707
 708 /* Public */
 709 void
 710 ArchiveEntry(Archive *AHX,
 711                          CatalogId catalogId, DumpId dumpId,
 712                          const char *tag,
 713                          const char *namespace,
 714                          const char *tablespace,
 715                          const char *owner, bool withOids,
 716                          const char *desc, teSection section,
 717                          const char *defn,
 718                          const char *dropStmt, const char *copyStmt,
 719                          const DumpId *deps, int nDeps,
 720                          DataDumperPtr dumpFn, void *dumpArg)
 721 {
 722         ArchiveHandle *AH = (ArchiveHandle *) AHX;
 723         TocEntry   *newToc;
 724
 725         newToc = (TocEntry *) calloc(1, sizeof(TocEntry));
 726         if (!newToc)
 727                 die_horribly(AH, modulename, "out of memory\n");
 728
 729         AH->tocCount++;
 730         if (dumpId > AH->maxDumpId)
 731                 AH->maxDumpId = dumpId;
 732
 733         newToc->prev = AH->toc->prev;
 734         newToc->next = AH->toc;
 735         AH->toc->prev->next = newToc;
 736         AH->toc->prev = newToc;
 737
 738         newToc->catalogId = catalogId;
 739         newToc->dumpId = dumpId;
 740         newToc->section = section;
 741
 742         newToc->tag = strdup(tag);
 743         newToc->namespace = namespace ? strdup(namespace) : NULL;
 744         newToc->tablespace = tablespace ? strdup(tablespace) : NULL;
 745         newToc->owner = strdup(owner);
 746         newToc->withOids = withOids;
 747         newToc->desc = strdup(desc);
 748         newToc->defn = strdup(defn);
 749         newToc->dropStmt = strdup(dropStmt);
 750         newToc->copyStmt = copyStmt ? strdup(copyStmt) : NULL;
 751
 752         if (nDeps > 0)
 753         {
 754                 newToc->dependencies = (DumpId *) malloc(nDeps * sizeof(DumpId));
 755                 memcpy(newToc->dependencies, deps, nDeps * sizeof(DumpId));
 756                 newToc->nDeps = nDeps;
 757         }
 758         else
 759         {
 760                 newToc->dependencies = NULL;
 761                 newToc->nDeps = 0;
 762         }
 763
 764         newToc->dataDumper = dumpFn;
 765         newToc->dataDumperArg = dumpArg;
 766         newToc->hadDumper = dumpFn ? true : false;
 767
 768         newToc->formatData = NULL;
 769
 770         if (AH->ArchiveEntryPtr !=NULL)
 771                 (*AH->ArchiveEntryPtr) (AH, newToc);
 772 }
 773
 774 /* Public */
 775 void
 776 PrintTOCSummary(Archive *AHX, RestoreOptions *ropt)
 777 {
 778         ArchiveHandle *AH = (ArchiveHandle *) AHX;
 779         TocEntry   *te;
 780         OutputContext sav;
 781         char       *fmtName;
 782
 783         if (ropt->filename)
 784                 sav = SetOutput(AH, ropt->filename, 0 /* no compression */ );
 785
 786         ahprintf(AH, ";\n; Archive created at %s", ctime(&AH->createDate));
 787         ahprintf(AH, ";     dbname: %s\n;     TOC Entries: %d\n;     Compression: %d\n",
 788                          AH->archdbname, AH->tocCount, AH->compression);
 789
 790         switch (AH->format)
 791         {
 792                 case archFiles:
 793                         fmtName = "FILES";
 794                         break;
 795                 case archCustom:
 796                         fmtName = "CUSTOM";
 797                         break;
 798                 case archTar:
 799                         fmtName = "TAR";
 800                         break;
 801                 default:
 802                         fmtName = "UNKNOWN";
 803         }
 804
 805         ahprintf(AH, ";     Dump Version: %d.%d-%d\n", AH->vmaj, AH->vmin, AH->vrev);
 806         ahprintf(AH, ";     Format: %s\n", fmtName);
 807         ahprintf(AH, ";     Integer: %d bytes\n", (int) AH->intSize);
 808         ahprintf(AH, ";     Offset: %d bytes\n", (int) AH->offSize);
 809         if (AH->archiveRemoteVersion)
 810                 ahprintf(AH, ";     Dumped from database version: %s\n",
 811                                  AH->archiveRemoteVersion);
 812         if (AH->archiveDumpVersion)
 813                 ahprintf(AH, ";     Dumped by pg_dump version: %s\n",
 814                                  AH->archiveDumpVersion);
 815
 816         ahprintf(AH, ";\n;\n; Selected TOC Entries:\n;\n");
 817
 818         /* We should print DATABASE entries whether or not -C was specified */
 819         ropt->createDB = 1;
 820
 821         for (te = AH->toc->next; te != AH->toc; te = te->next)
 822         {
 823                 if (ropt->verbose || _tocEntryRequired(te, ropt, true) != 0)
 824                         ahprintf(AH, "%d; %u %u %s %s %s %s\n", te->dumpId,
 825                                          te->catalogId.tableoid, te->catalogId.oid,
 826                                          te->desc, te->namespace ? te->namespace : "-",
 827                                          te->tag, te->owner);
 828                 if (ropt->verbose && te->nDeps > 0)
 829                 {
 830                         int                     i;
 831
 832                         ahprintf(AH, ";\tdepends on:");
 833                         for (i = 0; i < te->nDeps; i++)
 834                                 ahprintf(AH, " %d", te->dependencies[i]);
 835                         ahprintf(AH, "\n");
 836                 }
 837         }
 838
 839         if (ropt->filename)
 840                 ResetOutput(AH, sav);
 841 }
 842
 843 /***********
 844  * BLOB Archival
 845  ***********/
 846
 847 /* Called by a dumper to signal start of a BLOB */
 848 int
 849 StartBlob(Archive *AHX, Oid oid)
 850 {
 851         ArchiveHandle *AH = (ArchiveHandle *) AHX;
 852
 853         if (!AH->StartBlobPtr)
 854                 die_horribly(AH, modulename, "large-object output not supported in chosen format\n");
 855
 856         (*AH->StartBlobPtr) (AH, AH->currToc, oid);
 857
 858         return 1;
 859 }
 860
 861 /* Called by a dumper to signal end of a BLOB */
 862 int
 863 EndBlob(Archive *AHX, Oid oid)
 864 {
 865         ArchiveHandle *AH = (ArchiveHandle *) AHX;
 866
 867         if (AH->EndBlobPtr)
 868                 (*AH->EndBlobPtr) (AH, AH->currToc, oid);
 869
 870         return 1;
 871 }
 872
 873 /**********
 874  * BLOB Restoration
 875  **********/
 876
 877 /*
 878  * Called by a format handler before any blobs are restored
 879  */
 880 void
 881 StartRestoreBlobs(ArchiveHandle *AH)
 882 {
 883         if (!AH->ropt->single_txn)
 884         {
 885                 if (AH->connection)
 886                         StartTransaction(AH);
 887                 else
 888                         ahprintf(AH, "BEGIN;\n\n");
 889         }
 890
 891         AH->blobCount = 0;
 892 }
 893
 894 /*
 895  * Called by a format handler after all blobs are restored
 896  */
 897 void
 898 EndRestoreBlobs(ArchiveHandle *AH)
 899 {
 900         if (!AH->ropt->single_txn)
 901         {
 902                 if (AH->connection)
 903                         CommitTransaction(AH);
 904                 else
 905                         ahprintf(AH, "COMMIT;\n\n");
 906         }
 907
 908         ahlog(AH, 1, ngettext("restored %d large object\n",
 909                                                   "restored %d large objects\n",
 910                                                   AH->blobCount),
 911                   AH->blobCount);
 912 }
 913
 914
 915 /*
 916  * Called by a format handler to initiate restoration of a blob
 917  */
 918 void
 919 StartRestoreBlob(ArchiveHandle *AH, Oid oid, bool drop)
 920 {
 921         bool            old_blob_style = (AH->version < K_VERS_1_12);
 922         Oid                     loOid;
 923
 924         AH->blobCount++;
 925
 926         /* Initialize the LO Buffer */
 927         AH->lo_buf_used = 0;
 928
 929         ahlog(AH, 2, "restoring large object with OID %u\n", oid);
 930
 931         /* With an old archive we must do drop and create logic here */
 932         if (old_blob_style && drop)
 933                 DropBlobIfExists(AH, oid);
 934
 935         if (AH->connection)
 936         {
 937                 if (old_blob_style)
 938                 {
 939                         loOid = lo_create(AH->connection, oid);
 940                         if (loOid == 0 || loOid != oid)
 941                                 die_horribly(AH, modulename, "could not create large object %u: %s",
 942                                                          oid, PQerrorMessage(AH->connection));
 943                 }
 944                 AH->loFd = lo_open(AH->connection, oid, INV_WRITE);
 945                 if (AH->loFd == -1)
 946                         die_horribly(AH, modulename, "could not open large object %u: %s",
 947                                                  oid, PQerrorMessage(AH->connection));
 948         }
 949         else
 950         {
 951                 if (old_blob_style)
 952                         ahprintf(AH, "SELECT pg_catalog.lo_open(pg_catalog.lo_create('%u'), %d);\n",
 953                                          oid, INV_WRITE);
 954                 else
 955                         ahprintf(AH, "SELECT pg_catalog.lo_open('%u', %d);\n",
 956                                          oid, INV_WRITE);
 957         }
 958
 959         AH->writingBlob = 1;
 960 }
 961
 962 void
 963 EndRestoreBlob(ArchiveHandle *AH, Oid oid)
 964 {
 965         if (AH->lo_buf_used > 0)
 966         {
 967                 /* Write remaining bytes from the LO buffer */
 968                 dump_lo_buf(AH);
 969         }
 970
 971         AH->writingBlob = 0;
 972
 973         if (AH->connection)
 974         {
 975                 lo_close(AH->connection, AH->loFd);
 976                 AH->loFd = -1;
 977         }
 978         else
 979         {
 980                 ahprintf(AH, "SELECT pg_catalog.lo_close(0);\n\n");
 981         }
 982 }
 983
 984 /***********
 985  * Sorting and Reordering
 986  ***********/
 987
 988 void
 989 SortTocFromFile(Archive *AHX, RestoreOptions *ropt)
 990 {
 991         ArchiveHandle *AH = (ArchiveHandle *) AHX;
 992         FILE       *fh;
 993         char            buf[1024];
 994         char       *cmnt;
 995         char       *endptr;
 996         DumpId          id;
 997         TocEntry   *te;
 998         TocEntry   *tePrev;
 999
1000         /* Allocate space for the 'wanted' array, and init it */
1001         ropt->idWanted = (bool *) malloc(sizeof(bool) * AH->maxDumpId);
1002         memset(ropt->idWanted, 0, sizeof(bool) * AH->maxDumpId);
1003
1004         /* Set prev entry as head of list */
1005         tePrev = AH->toc;
1006
1007         /* Setup the file */
1008         fh = fopen(ropt->tocFile, PG_BINARY_R);
1009         if (!fh)
1010                 die_horribly(AH, modulename, "could not open TOC file \"%s\": %s\n",
1011                                          ropt->tocFile, strerror(errno));
1012
1013         while (fgets(buf, sizeof(buf), fh) != NULL)
1014         {
1015                 /* Truncate line at comment, if any */
1016                 cmnt = strchr(buf, ';');
1017                 if (cmnt != NULL)
1018                         cmnt[0] = '\0';
1019
1020                 /* Ignore if all blank */
1021                 if (strspn(buf, " \t\r") == strlen(buf))
1022                         continue;
1023
1024                 /* Get an ID, check it's valid and not already seen */
1025                 id = strtol(buf, &endptr, 10);
1026                 if (endptr == buf || id <= 0 || id > AH->maxDumpId ||
1027                         ropt->idWanted[id - 1])
1028                 {
1029                         write_msg(modulename, "WARNING: line ignored: %s\n", buf);
1030                         continue;
1031                 }
1032
1033                 /* Find TOC entry */
1034                 te = getTocEntryByDumpId(AH, id);
1035                 if (!te)
1036                         die_horribly(AH, modulename, "could not find entry for ID %d\n",
1037                                                  id);
1038
1039                 ropt->idWanted[id - 1] = true;
1040
1041                 _moveAfter(AH, tePrev, te);
1042                 tePrev = te;
1043         }
1044
1045         if (fclose(fh) != 0)
1046                 die_horribly(AH, modulename, "could not close TOC file: %s\n",
1047                                          strerror(errno));
1048 }
1049
1050 /*
1051  * Set up a dummy ID filter that selects all dump IDs
1052  */
1053 void
1054 InitDummyWantedList(Archive *AHX, RestoreOptions *ropt)
1055 {
1056         ArchiveHandle *AH = (ArchiveHandle *) AHX;
1057
1058         /* Allocate space for the 'wanted' array, and init it to 1's */
1059         ropt->idWanted = (bool *) malloc(sizeof(bool) * AH->maxDumpId);
1060         memset(ropt->idWanted, 1, sizeof(bool) * AH->maxDumpId);
1061 }
1062
1063 /**********************
1064  * 'Convenience functions that look like standard IO functions
1065  * for writing data when in dump mode.
1066  **********************/
1067
1068 /* Public */
1069 int
1070 archputs(const char *s, Archive *AH)
1071 {
1072         return WriteData(AH, s, strlen(s));
1073 }
1074
1075 /* Public */
1076 int
1077 archprintf(Archive *AH, const char *fmt,...)
1078 {
1079         char       *p = NULL;
1080         va_list         ap;
1081         int                     bSize = strlen(fmt) + 256;
1082         int                     cnt = -1;
1083
1084         /*
1085          * This is paranoid: deal with the possibility that vsnprintf is willing
1086          * to ignore trailing null or returns > 0 even if string does not fit. It
1087          * may be the case that it returns cnt = bufsize
1088          */
1089         while (cnt < 0 || cnt >= (bSize - 1))
1090         {
1091                 if (p != NULL)
1092                         free(p);
1093                 bSize *= 2;
1094                 p = (char *) malloc(bSize);
1095                 if (p == NULL)
1096                         exit_horribly(AH, modulename, "out of memory\n");
1097                 va_start(ap, fmt);
1098                 cnt = vsnprintf(p, bSize, fmt, ap);
1099                 va_end(ap);
1100         }
1101         WriteData(AH, p, cnt);
1102         free(p);
1103         return cnt;
1104 }
1105
1106
1107 /*******************************
1108  * Stuff below here should be 'private' to the archiver routines
1109  *******************************/
1110
1111 static OutputContext
1112 SetOutput(ArchiveHandle *AH, char *filename, int compression)
1113 {
1114         OutputContext sav;
1115         int                     fn;
1116
1117         /* Replace the AH output file handle */
1118         sav.OF = AH->OF;
1119         sav.gzOut = AH->gzOut;
1120
1121         if (filename)
1122                 fn = -1;
1123         else if (AH->FH)
1124                 fn = fileno(AH->FH);
1125         else if (AH->fSpec)
1126         {
1127                 fn = -1;
1128                 filename = AH->fSpec;
1129         }
1130         else
1131                 fn = fileno(stdout);
1132
1133         /* If compression explicitly requested, use gzopen */
1134 #ifdef HAVE_LIBZ
1135         if (compression != 0)
1136         {
1137                 char            fmode[10];
1138
1139                 /* Don't use PG_BINARY_x since this is zlib */
1140                 sprintf(fmode, "wb%d", compression);
1141                 if (fn >= 0)
1142                         AH->OF = gzdopen(dup(fn), fmode);
1143                 else
1144                         AH->OF = gzopen(filename, fmode);
1145                 AH->gzOut = 1;
1146         }
1147         else
1148 #endif
1149         {                                                       /* Use fopen */
1150                 if (AH->mode == archModeAppend)
1151                 {
1152                         if (fn >= 0)
1153                                 AH->OF = fdopen(dup(fn), PG_BINARY_A);
1154                         else
1155                                 AH->OF = fopen(filename, PG_BINARY_A);
1156                 }
1157                 else
1158                 {
1159                         if (fn >= 0)
1160                                 AH->OF = fdopen(dup(fn), PG_BINARY_W);
1161                         else
1162                                 AH->OF = fopen(filename, PG_BINARY_W);
1163                 }
1164                 AH->gzOut = 0;
1165         }
1166
1167         if (!AH->OF)
1168         {
1169                 if (filename)
1170                         die_horribly(AH, modulename, "could not open output file \"%s\": %s\n",
1171                                                  filename, strerror(errno));
1172                 else
1173                         die_horribly(AH, modulename, "could not open output file: %s\n",
1174                                                  strerror(errno));
1175         }
1176
1177         return sav;
1178 }
1179
1180 static void
1181 ResetOutput(ArchiveHandle *AH, OutputContext sav)
1182 {
1183         int                     res;
1184
1185         if (AH->gzOut)
1186                 res = GZCLOSE(AH->OF);
1187         else
1188                 res = fclose(AH->OF);
1189
1190         if (res != 0)
1191                 die_horribly(AH, modulename, "could not close output file: %s\n",
1192                                          strerror(errno));
1193
1194         AH->gzOut = sav.gzOut;
1195         AH->OF = sav.OF;
1196 }
1197
1198
1199
1200 /*
1201  *      Print formatted text to the output file (usually stdout).
1202  */
1203 int
1204 ahprintf(ArchiveHandle *AH, const char *fmt,...)
1205 {
1206         char       *p = NULL;
1207         va_list         ap;
1208         int                     bSize = strlen(fmt) + 256;              /* Should be enough */
1209         int                     cnt = -1;
1210
1211         /*
1212          * This is paranoid: deal with the possibility that vsnprintf is willing
1213          * to ignore trailing null
1214          */
1215
1216         /*
1217          * or returns > 0 even if string does not fit. It may be the case that it
1218          * returns cnt = bufsize
1219          */
1220         while (cnt < 0 || cnt >= (bSize - 1))
1221         {
1222                 if (p != NULL)
1223                         free(p);
1224                 bSize *= 2;
1225                 p = (char *) malloc(bSize);
1226                 if (p == NULL)
1227                         die_horribly(AH, modulename, "out of memory\n");
1228                 va_start(ap, fmt);
1229                 cnt = vsnprintf(p, bSize, fmt, ap);
1230                 va_end(ap);
1231         }
1232         ahwrite(p, 1, cnt, AH);
1233         free(p);
1234         return cnt;
1235 }
1236
1237 void
1238 ahlog(ArchiveHandle *AH, int level, const char *fmt,...)
1239 {
1240         va_list         ap;
1241
1242         if (AH->debugLevel < level && (!AH->public.verbose || level > 1))
1243                 return;
1244
1245         va_start(ap, fmt);
1246         _write_msg(NULL, fmt, ap);
1247         va_end(ap);
1248 }
1249
1250 /*
1251  * Single place for logic which says 'We are restoring to a direct DB connection'.
1252  */
1253 static int
1254 RestoringToDB(ArchiveHandle *AH)
1255 {
1256         return (AH->ropt && AH->ropt->useDB && AH->connection);
1257 }
1258
1259 /*
1260  * Dump the current contents of the LO data buffer while writing a BLOB
1261  */
1262 static void
1263 dump_lo_buf(ArchiveHandle *AH)
1264 {
1265         if (AH->connection)
1266         {
1267                 size_t          res;
1268
1269                 res = lo_write(AH->connection, AH->loFd, AH->lo_buf, AH->lo_buf_used);
1270                 ahlog(AH, 5, ngettext("wrote %lu byte of large object data (result = %lu)\n",
1271                                          "wrote %lu bytes of large object data (result = %lu)\n",
1272                                                           AH->lo_buf_used),
1273                           (unsigned long) AH->lo_buf_used, (unsigned long) res);
1274                 if (res != AH->lo_buf_used)
1275                         die_horribly(AH, modulename,
1276                         "could not write to large object (result: %lu, expected: %lu)\n",
1277                                            (unsigned long) res, (unsigned long) AH->lo_buf_used);
1278         }
1279         else
1280         {
1281                 PQExpBuffer buf = createPQExpBuffer();
1282
1283                 appendByteaLiteralAHX(buf,
1284                                                           (const unsigned char *) AH->lo_buf,
1285                                                           AH->lo_buf_used,
1286                                                           AH);
1287
1288                 /* Hack: turn off writingBlob so ahwrite doesn't recurse to here */
1289                 AH->writingBlob = 0;
1290                 ahprintf(AH, "SELECT pg_catalog.lowrite(0, %s);\n", buf->data);
1291                 AH->writingBlob = 1;
1292
1293                 destroyPQExpBuffer(buf);
1294         }
1295         AH->lo_buf_used = 0;
1296 }
1297
1298
1299 /*
1300  *      Write buffer to the output file (usually stdout). This is user for
1301  *      outputting 'restore' scripts etc. It is even possible for an archive
1302  *      format to create a custom output routine to 'fake' a restore if it
1303  *      wants to generate a script (see TAR output).
1304  */
1305 int
1306 ahwrite(const void *ptr, size_t size, size_t nmemb, ArchiveHandle *AH)
1307 {
1308         size_t          res;
1309
1310         if (AH->writingBlob)
1311         {
1312                 size_t          remaining = size * nmemb;
1313
1314                 while (AH->lo_buf_used + remaining > AH->lo_buf_size)
1315                 {
1316                         size_t          avail = AH->lo_buf_size - AH->lo_buf_used;
1317
1318                         memcpy((char *) AH->lo_buf + AH->lo_buf_used, ptr, avail);
1319                         ptr = (const void *) ((const char *) ptr + avail);
1320                         remaining -= avail;
1321                         AH->lo_buf_used += avail;
1322                         dump_lo_buf(AH);
1323                 }
1324
1325                 memcpy((char *) AH->lo_buf + AH->lo_buf_used, ptr, remaining);
1326                 AH->lo_buf_used += remaining;
1327
1328                 return size * nmemb;
1329         }
1330         else if (AH->gzOut)
1331         {
1332                 res = GZWRITE((void *) ptr, size, nmemb, AH->OF);
1333                 if (res != (nmemb * size))
1334                         die_horribly(AH, modulename, "could not write to output file: %s\n", strerror(errno));
1335                 return res;
1336         }
1337         else if (AH->CustomOutPtr)
1338         {
1339                 res = AH->CustomOutPtr (AH, ptr, size * nmemb);
1340
1341                 if (res != (nmemb * size))
1342                         die_horribly(AH, modulename, "could not write to custom output routine\n");
1343                 return res;
1344         }
1345         else
1346         {
1347                 /*
1348                  * If we're doing a restore, and it's direct to DB, and we're
1349                  * connected then send it to the DB.
1350                  */
1351                 if (RestoringToDB(AH))
1352                         return ExecuteSqlCommandBuf(AH, (void *) ptr, size * nmemb);            /* Always 1, currently */
1353                 else
1354                 {
1355                         res = fwrite((void *) ptr, size, nmemb, AH->OF);
1356                         if (res != nmemb)
1357                                 die_horribly(AH, modulename, "could not write to output file: %s\n",
1358                                                          strerror(errno));
1359                         return res;
1360                 }
1361         }
1362 }
1363
1364 /* Common exit code */
1365 static void
1366 _write_msg(const char *modulename, const char *fmt, va_list ap)
1367 {
1368         if (modulename)
1369                 fprintf(stderr, "%s: [%s] ", progname, _(modulename));
1370         else
1371                 fprintf(stderr, "%s: ", progname);
1372         vfprintf(stderr, _(fmt), ap);
1373 }
1374
1375 void
1376 write_msg(const char *modulename, const char *fmt,...)
1377 {
1378         va_list         ap;
1379
1380         va_start(ap, fmt);
1381         _write_msg(modulename, fmt, ap);
1382         va_end(ap);
1383 }
1384
1385
1386 static void
1387 _die_horribly(ArchiveHandle *AH, const char *modulename, const char *fmt, va_list ap)
1388 {
1389         _write_msg(modulename, fmt, ap);
1390
1391         if (AH)
1392         {
1393                 if (AH->public.verbose)
1394                         write_msg(NULL, "*** aborted because of error\n");
1395                 if (AH->connection)
1396                         PQfinish(AH->connection);
1397         }
1398
1399         exit(1);
1400 }
1401
1402 /* External use */
1403 void
1404 exit_horribly(Archive *AH, const char *modulename, const char *fmt,...)
1405 {
1406         va_list         ap;
1407
1408         va_start(ap, fmt);
1409         _die_horribly((ArchiveHandle *) AH, modulename, fmt, ap);
1410         va_end(ap);
1411 }
1412
1413 /* Archiver use (just different arg declaration) */
1414 void
1415 die_horribly(ArchiveHandle *AH, const char *modulename, const char *fmt,...)
1416 {
1417         va_list         ap;
1418
1419         va_start(ap, fmt);
1420         _die_horribly(AH, modulename, fmt, ap);
1421         va_end(ap);
1422 }
1423
1424 /* on some error, we may decide to go on... */
1425 void
1426 warn_or_die_horribly(ArchiveHandle *AH,
1427                                          const char *modulename, const char *fmt,...)
1428 {
1429         va_list         ap;
1430
1431         switch (AH->stage)
1432         {
1433
1434                 case STAGE_NONE:
1435                         /* Do nothing special */
1436                         break;
1437
1438                 case STAGE_INITIALIZING:
1439                         if (AH->stage != AH->lastErrorStage)
1440                                 write_msg(modulename, "Error while INITIALIZING:\n");
1441                         break;
1442
1443                 case STAGE_PROCESSING:
1444                         if (AH->stage != AH->lastErrorStage)
1445                                 write_msg(modulename, "Error while PROCESSING TOC:\n");
1446                         break;
1447
1448                 case STAGE_FINALIZING:
1449                         if (AH->stage != AH->lastErrorStage)
1450                                 write_msg(modulename, "Error while FINALIZING:\n");
1451                         break;
1452         }
1453         if (AH->currentTE != NULL && AH->currentTE != AH->lastErrorTE)
1454         {
1455                 write_msg(modulename, "Error from TOC entry %d; %u %u %s %s %s\n",
1456                                   AH->currentTE->dumpId,
1457                          AH->currentTE->catalogId.tableoid, AH->currentTE->catalogId.oid,
1458                           AH->currentTE->desc, AH->currentTE->tag, AH->currentTE->owner);
1459         }
1460         AH->lastErrorStage = AH->stage;
1461         AH->lastErrorTE = AH->currentTE;
1462
1463         va_start(ap, fmt);
1464         if (AH->public.exit_on_error)
1465                 _die_horribly(AH, modulename, fmt, ap);
1466         else
1467         {
1468                 _write_msg(modulename, fmt, ap);
1469                 AH->public.n_errors++;
1470         }
1471         va_end(ap);
1472 }
1473
1474 static void
1475 _moveAfter(ArchiveHandle *AH, TocEntry *pos, TocEntry *te)
1476 {
1477         te->prev->next = te->next;
1478         te->next->prev = te->prev;
1479
1480         te->prev = pos;
1481         te->next = pos->next;
1482
1483         pos->next->prev = te;
1484         pos->next = te;
1485 }
1486
1487 #ifdef NOT_USED
1488
1489 static void
1490 _moveBefore(ArchiveHandle *AH, TocEntry *pos, TocEntry *te)
1491 {
1492         te->prev->next = te->next;
1493         te->next->prev = te->prev;
1494
1495         te->prev = pos->prev;
1496         te->next = pos;
1497         pos->prev->next = te;
1498         pos->prev = te;
1499 }
1500 #endif
1501
1502 static TocEntry *
1503 getTocEntryByDumpId(ArchiveHandle *AH, DumpId id)
1504 {
1505         TocEntry   *te;
1506
1507         for (te = AH->toc->next; te != AH->toc; te = te->next)
1508         {
1509                 if (te->dumpId == id)
1510                         return te;
1511         }
1512         return NULL;
1513 }
1514
1515 teReqs
1516 TocIDRequired(ArchiveHandle *AH, DumpId id, RestoreOptions *ropt)
1517 {
1518         TocEntry   *te = getTocEntryByDumpId(AH, id);
1519
1520         if (!te)
1521                 return 0;
1522
1523         return _tocEntryRequired(te, ropt, true);
1524 }
1525
1526 size_t
1527 WriteOffset(ArchiveHandle *AH, pgoff_t o, int wasSet)
1528 {
1529         int                     off;
1530
1531         /* Save the flag */
1532         (*AH->WriteBytePtr) (AH, wasSet);
1533
1534         /* Write out pgoff_t smallest byte first, prevents endian mismatch */
1535         for (off = 0; off < sizeof(pgoff_t); off++)
1536         {
1537                 (*AH->WriteBytePtr) (AH, o & 0xFF);
1538                 o >>= 8;
1539         }
1540         return sizeof(pgoff_t) + 1;
1541 }
1542
1543 int
1544 ReadOffset(ArchiveHandle *AH, pgoff_t * o)
1545 {
1546         int                     i;
1547         int                     off;
1548         int                     offsetFlg;
1549
1550         /* Initialize to zero */
1551         *o = 0;
1552
1553         /* Check for old version */
1554         if (AH->version < K_VERS_1_7)
1555         {
1556                 /* Prior versions wrote offsets using WriteInt */
1557                 i = ReadInt(AH);
1558                 /* -1 means not set */
1559                 if (i < 0)
1560                         return K_OFFSET_POS_NOT_SET;
1561                 else if (i == 0)
1562                         return K_OFFSET_NO_DATA;
1563
1564                 /* Cast to pgoff_t because it was written as an int. */
1565                 *o = (pgoff_t) i;
1566                 return K_OFFSET_POS_SET;
1567         }
1568
1569         /*
1570          * Read the flag indicating the state of the data pointer. Check if valid
1571          * and die if not.
1572          *
1573          * This used to be handled by a negative or zero pointer, now we use an
1574          * extra byte specifically for the state.
1575          */
1576         offsetFlg = (*AH->ReadBytePtr) (AH) & 0xFF;
1577
1578         switch (offsetFlg)
1579         {
1580                 case K_OFFSET_POS_NOT_SET:
1581                 case K_OFFSET_NO_DATA:
1582                 case K_OFFSET_POS_SET:
1583
1584                         break;
1585
1586                 default:
1587                         die_horribly(AH, modulename, "unexpected data offset flag %d\n", offsetFlg);
1588         }
1589
1590         /*
1591          * Read the bytes
1592          */
1593         for (off = 0; off < AH->offSize; off++)
1594         {
1595                 if (off < sizeof(pgoff_t))
1596                         *o |= ((pgoff_t) ((*AH->ReadBytePtr) (AH))) << (off * 8);
1597                 else
1598                 {
1599                         if ((*AH->ReadBytePtr) (AH) != 0)
1600                                 die_horribly(AH, modulename, "file offset in dump file is too large\n");
1601                 }
1602         }
1603
1604         return offsetFlg;
1605 }
1606
1607 size_t
1608 WriteInt(ArchiveHandle *AH, int i)
1609 {
1610         int                     b;
1611
1612         /*
1613          * This is a bit yucky, but I don't want to make the binary format very
1614          * dependent on representation, and not knowing much about it, I write out
1615          * a sign byte. If you change this, don't forget to change the file
1616          * version #, and modify readInt to read the new format AS WELL AS the old
1617          * formats.
1618          */
1619
1620         /* SIGN byte */
1621         if (i < 0)
1622         {
1623                 (*AH->WriteBytePtr) (AH, 1);
1624                 i = -i;
1625         }
1626         else
1627                 (*AH->WriteBytePtr) (AH, 0);
1628
1629         for (b = 0; b < AH->intSize; b++)
1630         {
1631                 (*AH->WriteBytePtr) (AH, i & 0xFF);
1632                 i >>= 8;
1633         }
1634
1635         return AH->intSize + 1;
1636 }
1637
1638 int
1639 ReadInt(ArchiveHandle *AH)
1640 {
1641         int                     res = 0;
1642         int                     bv,
1643                                 b;
1644         int                     sign = 0;               /* Default positive */
1645         int                     bitShift = 0;
1646
1647         if (AH->version > K_VERS_1_0)
1648                 /* Read a sign byte */
1649                 sign = (*AH->ReadBytePtr) (AH);
1650
1651         for (b = 0; b < AH->intSize; b++)
1652         {
1653                 bv = (*AH->ReadBytePtr) (AH) & 0xFF;
1654                 if (bv != 0)
1655                         res = res + (bv << bitShift);
1656                 bitShift += 8;
1657         }
1658
1659         if (sign)
1660                 res = -res;
1661
1662         return res;
1663 }
1664
1665 size_t
1666 WriteStr(ArchiveHandle *AH, const char *c)
1667 {
1668         size_t          res;
1669
1670         if (c)
1671         {
1672                 res = WriteInt(AH, strlen(c));
1673                 res += (*AH->WriteBufPtr) (AH, c, strlen(c));
1674         }
1675         else
1676                 res = WriteInt(AH, -1);
1677
1678         return res;
1679 }
1680
1681 char *
1682 ReadStr(ArchiveHandle *AH)
1683 {
1684         char       *buf;
1685         int                     l;
1686
1687         l = ReadInt(AH);
1688         if (l < 0)
1689                 buf = NULL;
1690         else
1691         {
1692                 buf = (char *) malloc(l + 1);
1693                 if (!buf)
1694                         die_horribly(AH, modulename, "out of memory\n");
1695
1696                 if ((*AH->ReadBufPtr) (AH, (void *) buf, l) != l)
1697                         die_horribly(AH, modulename, "unexpected end of file\n");
1698
1699                 buf[l] = '\0';
1700         }
1701
1702         return buf;
1703 }
1704
1705 static int
1706 _discoverArchiveFormat(ArchiveHandle *AH)
1707 {
1708         FILE       *fh;
1709         char            sig[6];                 /* More than enough */
1710         size_t          cnt;
1711         int                     wantClose = 0;
1712
1713 #if 0
1714         write_msg(modulename, "attempting to ascertain archive format\n");
1715 #endif
1716
1717         if (AH->lookahead)
1718                 free(AH->lookahead);
1719
1720         AH->lookaheadSize = 512;
1721         AH->lookahead = calloc(1, 512);
1722         AH->lookaheadLen = 0;
1723         AH->lookaheadPos = 0;
1724
1725         if (AH->fSpec)
1726         {
1727                 wantClose = 1;
1728                 fh = fopen(AH->fSpec, PG_BINARY_R);
1729                 if (!fh)
1730                         die_horribly(AH, modulename, "could not open input file \"%s\": %s\n",
1731                                                  AH->fSpec, strerror(errno));
1732         }
1733         else
1734         {
1735                 fh = stdin;
1736                 if (!fh)
1737                         die_horribly(AH, modulename, "could not open input file: %s\n",
1738                                                  strerror(errno));
1739         }
1740
1741         cnt = fread(sig, 1, 5, fh);
1742
1743         if (cnt != 5)
1744         {
1745                 if (ferror(fh))
1746                         die_horribly(AH, modulename, "could not read input file: %s\n", strerror(errno));
1747                 else
1748                         die_horribly(AH, modulename, "input file is too short (read %lu, expected 5)\n",
1749                                                  (unsigned long) cnt);
1750         }
1751
1752         /* Save it, just in case we need it later */
1753         strncpy(&AH->lookahead[0], sig, 5);
1754         AH->lookaheadLen = 5;
1755
1756         if (strncmp(sig, "PGDMP", 5) == 0)
1757         {
1758                 /*
1759                  * Finish reading (most of) a custom-format header.
1760                  *
1761                  * NB: this code must agree with ReadHead().
1762                  */
1763                 AH->vmaj = fgetc(fh);
1764                 AH->vmin = fgetc(fh);
1765
1766                 /* Save these too... */
1767                 AH->lookahead[AH->lookaheadLen++] = AH->vmaj;
1768                 AH->lookahead[AH->lookaheadLen++] = AH->vmin;
1769
1770                 /* Check header version; varies from V1.0 */
1771                 if (AH->vmaj > 1 || ((AH->vmaj == 1) && (AH->vmin > 0)))                /* Version > 1.0 */
1772                 {
1773                         AH->vrev = fgetc(fh);
1774                         AH->lookahead[AH->lookaheadLen++] = AH->vrev;
1775                 }
1776                 else
1777                         AH->vrev = 0;
1778
1779                 /* Make a convenient integer <maj><min><rev>00 */
1780                 AH->version = ((AH->vmaj * 256 + AH->vmin) * 256 + AH->vrev) * 256 + 0;
1781
1782                 AH->intSize = fgetc(fh);
1783                 AH->lookahead[AH->lookaheadLen++] = AH->intSize;
1784
1785                 if (AH->version >= K_VERS_1_7)
1786                 {
1787                         AH->offSize = fgetc(fh);
1788                         AH->lookahead[AH->lookaheadLen++] = AH->offSize;
1789                 }
1790                 else
1791                         AH->offSize = AH->intSize;
1792
1793                 AH->format = fgetc(fh);
1794                 AH->lookahead[AH->lookaheadLen++] = AH->format;
1795         }
1796         else
1797         {
1798                 /*
1799                  * *Maybe* we have a tar archive format file... So, read first 512
1800                  * byte header...
1801                  */
1802                 cnt = fread(&AH->lookahead[AH->lookaheadLen], 1, 512 - AH->lookaheadLen, fh);
1803                 AH->lookaheadLen += cnt;
1804
1805                 if (AH->lookaheadLen != 512)
1806                         die_horribly(AH, modulename, "input file does not appear to be a valid archive (too short?)\n");
1807
1808                 if (!isValidTarHeader(AH->lookahead))
1809                         die_horribly(AH, modulename, "input file does not appear to be a valid archive\n");
1810
1811                 AH->format = archTar;
1812         }
1813
1814         /* If we can't seek, then mark the header as read */
1815         if (fseeko(fh, 0, SEEK_SET) != 0)
1816         {
1817                 /*
1818                  * NOTE: Formats that use the lookahead buffer can unset this in their
1819                  * Init routine.
1820                  */
1821                 AH->readHeader = 1;
1822         }
1823         else
1824                 AH->lookaheadLen = 0;   /* Don't bother since we've reset the file */
1825
1826         /* Close the file */
1827         if (wantClose)
1828                 if (fclose(fh) != 0)
1829                         die_horribly(AH, modulename, "could not close input file: %s\n",
1830                                                  strerror(errno));
1831
1832         return AH->format;
1833 }
1834
1835
1836 /*
1837  * Allocate an archive handle
1838  */
1839 static ArchiveHandle *
1840 _allocAH(const char *FileSpec, const ArchiveFormat fmt,
1841                  const int compression, ArchiveMode mode)
1842 {
1843         ArchiveHandle *AH;
1844
1845 #if 0
1846         write_msg(modulename, "allocating AH for %s, format %d\n", FileSpec, fmt);
1847 #endif
1848
1849         AH = (ArchiveHandle *) calloc(1, sizeof(ArchiveHandle));
1850         if (!AH)
1851                 die_horribly(AH, modulename, "out of memory\n");
1852
1853         /* AH->debugLevel = 100; */
1854
1855         AH->vmaj = K_VERS_MAJOR;
1856         AH->vmin = K_VERS_MINOR;
1857         AH->vrev = K_VERS_REV;
1858
1859         /* Make a convenient integer <maj><min><rev>00 */
1860         AH->version = ((AH->vmaj * 256 + AH->vmin) * 256 + AH->vrev) * 256 + 0;
1861
1862         /* initialize for backwards compatible string processing */
1863         AH->public.encoding = 0;        /* PG_SQL_ASCII */
1864         AH->public.std_strings = false;
1865
1866         /* sql error handling */
1867         AH->public.exit_on_error = true;
1868         AH->public.n_errors = 0;
1869
1870         AH->archiveDumpVersion = PG_VERSION;
1871
1872         AH->createDate = time(NULL);
1873
1874         AH->intSize = sizeof(int);
1875         AH->offSize = sizeof(pgoff_t);
1876         if (FileSpec)
1877         {
1878                 AH->fSpec = strdup(FileSpec);
1879
1880                 /*
1881                  * Not used; maybe later....
1882                  *
1883                  * AH->workDir = strdup(FileSpec); for(i=strlen(FileSpec) ; i > 0 ;
1884                  * i--) if (AH->workDir[i-1] == '/')
1885                  */
1886         }
1887         else
1888                 AH->fSpec = NULL;
1889
1890         AH->currUser = NULL;            /* unknown */
1891         AH->currSchema = NULL;          /* ditto */
1892         AH->currTablespace = NULL;      /* ditto */
1893         AH->currWithOids = -1;          /* force SET */
1894
1895         AH->toc = (TocEntry *) calloc(1, sizeof(TocEntry));
1896         if (!AH->toc)
1897                 die_horribly(AH, modulename, "out of memory\n");
1898
1899         AH->toc->next = AH->toc;
1900         AH->toc->prev = AH->toc;
1901
1902         AH->mode = mode;
1903         AH->compression = compression;
1904
1905         AH->pgCopyBuf = createPQExpBuffer();
1906         AH->sqlBuf = createPQExpBuffer();
1907
1908         /* Open stdout with no compression for AH output handle */
1909         AH->gzOut = 0;
1910         AH->OF = stdout;
1911
1912         /*
1913          * On Windows, we need to use binary mode to read/write non-text archive
1914          * formats.  Force stdin/stdout into binary mode if that is what we are
1915          * using.
1916          */
1917 #ifdef WIN32
1918         if (fmt != archNull &&
1919                 (AH->fSpec == NULL || strcmp(AH->fSpec, "") == 0))
1920         {
1921                 if (mode == archModeWrite)
1922                         setmode(fileno(stdout), O_BINARY);
1923                 else
1924                         setmode(fileno(stdin), O_BINARY);
1925         }
1926 #endif
1927
1928         if (fmt == archUnknown)
1929                 AH->format = _discoverArchiveFormat(AH);
1930         else
1931                 AH->format = fmt;
1932
1933         AH->promptPassword = TRI_DEFAULT;
1934
1935         switch (AH->format)
1936         {
1937                 case archCustom:
1938                         InitArchiveFmt_Custom(AH);
1939                         break;
1940
1941                 case archFiles:
1942                         InitArchiveFmt_Files(AH);
1943                         break;
1944
1945                 case archNull:
1946                         InitArchiveFmt_Null(AH);
1947                         break;
1948
1949                 case archTar:
1950                         InitArchiveFmt_Tar(AH);
1951                         break;
1952
1953                 default:
1954                         die_horribly(AH, modulename, "unrecognized file format \"%d\"\n", fmt);
1955         }
1956
1957         return AH;
1958 }
1959
1960
1961 void
1962 WriteDataChunks(ArchiveHandle *AH)
1963 {
1964         TocEntry   *te;
1965         StartDataPtr startPtr;
1966         EndDataPtr      endPtr;
1967
1968         for (te = AH->toc->next; te != AH->toc; te = te->next)
1969         {
1970                 if (te->dataDumper != NULL)
1971                 {
1972                         AH->currToc = te;
1973                         /* printf("Writing data for %d (%x)\n", te->id, te); */
1974
1975                         if (strcmp(te->desc, "BLOBS") == 0)
1976                         {
1977                                 startPtr = AH->StartBlobsPtr;
1978                                 endPtr = AH->EndBlobsPtr;
1979                         }
1980                         else
1981                         {
1982                                 startPtr = AH->StartDataPtr;
1983                                 endPtr = AH->EndDataPtr;
1984                         }
1985
1986                         if (startPtr != NULL)
1987                                 (*startPtr) (AH, te);
1988
1989                         /*
1990                          * printf("Dumper arg for %d is %x\n", te->id, te->dataDumperArg);
1991                          */
1992
1993                         /*
1994                          * The user-provided DataDumper routine needs to call
1995                          * AH->WriteData
1996                          */
1997                         (*te->dataDumper) ((Archive *) AH, te->dataDumperArg);
1998
1999                         if (endPtr != NULL)
2000                                 (*endPtr) (AH, te);
2001                         AH->currToc = NULL;
2002                 }
2003         }
2004 }
2005
2006 void
2007 WriteToc(ArchiveHandle *AH)
2008 {
2009         TocEntry   *te;
2010         char            workbuf[32];
2011         int                     i;
2012
2013         /* printf("%d TOC Entries to save\n", AH->tocCount); */
2014
2015         WriteInt(AH, AH->tocCount);
2016
2017         for (te = AH->toc->next; te != AH->toc; te = te->next)
2018         {
2019                 WriteInt(AH, te->dumpId);
2020                 WriteInt(AH, te->dataDumper ? 1 : 0);
2021
2022                 /* OID is recorded as a string for historical reasons */
2023                 sprintf(workbuf, "%u", te->catalogId.tableoid);
2024                 WriteStr(AH, workbuf);
2025                 sprintf(workbuf, "%u", te->catalogId.oid);
2026                 WriteStr(AH, workbuf);
2027
2028                 WriteStr(AH, te->tag);
2029                 WriteStr(AH, te->desc);
2030                 WriteInt(AH, te->section);
2031                 WriteStr(AH, te->defn);
2032                 WriteStr(AH, te->dropStmt);
2033                 WriteStr(AH, te->copyStmt);
2034                 WriteStr(AH, te->namespace);
2035                 WriteStr(AH, te->tablespace);
2036                 WriteStr(AH, te->owner);
2037                 WriteStr(AH, te->withOids ? "true" : "false");
2038
2039                 /* Dump list of dependencies */
2040                 for (i = 0; i < te->nDeps; i++)
2041                 {
2042                         sprintf(workbuf, "%d", te->dependencies[i]);
2043                         WriteStr(AH, workbuf);
2044                 }
2045                 WriteStr(AH, NULL);             /* Terminate List */
2046
2047                 if (AH->WriteExtraTocPtr)
2048                         (*AH->WriteExtraTocPtr) (AH, te);
2049         }
2050 }
2051
2052 void
2053 ReadToc(ArchiveHandle *AH)
2054 {
2055         int                     i;
2056         char       *tmp;
2057         DumpId     *deps;
2058         int                     depIdx;
2059         int                     depSize;
2060         TocEntry   *te;
2061
2062         AH->tocCount = ReadInt(AH);
2063         AH->maxDumpId = 0;
2064
2065         for (i = 0; i < AH->tocCount; i++)
2066         {
2067                 te = (TocEntry *) calloc(1, sizeof(TocEntry));
2068                 te->dumpId = ReadInt(AH);
2069
2070                 if (te->dumpId > AH->maxDumpId)
2071                         AH->maxDumpId = te->dumpId;
2072
2073                 /* Sanity check */
2074                 if (te->dumpId <= 0)
2075                         die_horribly(AH, modulename,
2076                                            "entry ID %d out of range -- perhaps a corrupt TOC\n",
2077                                                  te->dumpId);
2078
2079                 te->hadDumper = ReadInt(AH);
2080
2081                 if (AH->version >= K_VERS_1_8)
2082                 {
2083                         tmp = ReadStr(AH);
2084                         sscanf(tmp, "%u", &te->catalogId.tableoid);
2085                         free(tmp);
2086                 }
2087                 else
2088                         te->catalogId.tableoid = InvalidOid;
2089                 tmp = ReadStr(AH);
2090                 sscanf(tmp, "%u", &te->catalogId.oid);
2091                 free(tmp);
2092
2093                 te->tag = ReadStr(AH);
2094                 te->desc = ReadStr(AH);
2095
2096                 if (AH->version >= K_VERS_1_11)
2097                 {
2098                         te->section = ReadInt(AH);
2099                 }
2100                 else
2101                 {
2102                         /*
2103                          * Rules for pre-8.4 archives wherein pg_dump hasn't classified
2104                          * the entries into sections.  This list need not cover entry
2105                          * types added later than 8.4.
2106                          */
2107                         if (strcmp(te->desc, "COMMENT") == 0 ||
2108                                 strcmp(te->desc, "ACL") == 0 ||
2109                                 strcmp(te->desc, "ACL LANGUAGE") == 0)
2110                                 te->section = SECTION_NONE;
2111                         else if (strcmp(te->desc, "TABLE DATA") == 0 ||
2112                                          strcmp(te->desc, "BLOBS") == 0 ||
2113                                          strcmp(te->desc, "BLOB COMMENTS") == 0)
2114                                 te->section = SECTION_DATA;
2115                         else if (strcmp(te->desc, "CONSTRAINT") == 0 ||
2116                                          strcmp(te->desc, "CHECK CONSTRAINT") == 0 ||
2117                                          strcmp(te->desc, "FK CONSTRAINT") == 0 ||
2118                                          strcmp(te->desc, "INDEX") == 0 ||
2119                                          strcmp(te->desc, "RULE") == 0 ||
2120                                          strcmp(te->desc, "TRIGGER") == 0)
2121                                 te->section = SECTION_POST_DATA;
2122                         else
2123                                 te->section = SECTION_PRE_DATA;
2124                 }
2125
2126                 te->defn = ReadStr(AH);
2127                 te->dropStmt = ReadStr(AH);
2128
2129                 if (AH->version >= K_VERS_1_3)
2130                         te->copyStmt = ReadStr(AH);
2131
2132                 if (AH->version >= K_VERS_1_6)
2133                         te->namespace = ReadStr(AH);
2134
2135                 if (AH->version >= K_VERS_1_10)
2136                         te->tablespace = ReadStr(AH);
2137
2138                 te->owner = ReadStr(AH);
2139                 if (AH->version >= K_VERS_1_9)
2140                 {
2141                         if (strcmp(ReadStr(AH), "true") == 0)
2142                                 te->withOids = true;
2143                         else
2144                                 te->withOids = false;
2145                 }
2146                 else
2147                         te->withOids = true;
2148
2149                 /* Read TOC entry dependencies */
2150                 if (AH->version >= K_VERS_1_5)
2151                 {
2152                         depSize = 100;
2153                         deps = (DumpId *) malloc(sizeof(DumpId) * depSize);
2154                         depIdx = 0;
2155                         for (;;)
2156                         {
2157                                 tmp = ReadStr(AH);
2158                                 if (!tmp)
2159                                         break;          /* end of list */
2160                                 if (depIdx >= depSize)
2161                                 {
2162                                         depSize *= 2;
2163                                         deps = (DumpId *) realloc(deps, sizeof(DumpId) * depSize);
2164                                 }
2165                                 sscanf(tmp, "%d", &deps[depIdx]);
2166                                 free(tmp);
2167                                 depIdx++;
2168                         }
2169
2170                         if (depIdx > 0)         /* We have a non-null entry */
2171                         {
2172                                 deps = (DumpId *) realloc(deps, sizeof(DumpId) * depIdx);
2173                                 te->dependencies = deps;
2174                                 te->nDeps = depIdx;
2175                         }
2176                         else
2177                         {
2178                                 free(deps);
2179                                 te->dependencies = NULL;
2180                                 te->nDeps = 0;
2181                         }
2182                 }
2183                 else
2184                 {
2185                         te->dependencies = NULL;
2186                         te->nDeps = 0;
2187                 }
2188
2189                 if (AH->ReadExtraTocPtr)
2190                         (*AH->ReadExtraTocPtr) (AH, te);
2191
2192                 ahlog(AH, 3, "read TOC entry %d (ID %d) for %s %s\n",
2193                           i, te->dumpId, te->desc, te->tag);
2194
2195                 /* link completed entry into TOC circular list */
2196                 te->prev = AH->toc->prev;
2197                 AH->toc->prev->next = te;
2198                 AH->toc->prev = te;
2199                 te->next = AH->toc;
2200
2201                 /* special processing immediately upon read for some items */
2202                 if (strcmp(te->desc, "ENCODING") == 0)
2203                         processEncodingEntry(AH, te);
2204                 else if (strcmp(te->desc, "STDSTRINGS") == 0)
2205                         processStdStringsEntry(AH, te);
2206         }
2207 }
2208
2209 static void
2210 processEncodingEntry(ArchiveHandle *AH, TocEntry *te)
2211 {
2212         /* te->defn should have the form SET client_encoding = 'foo'; */
2213         char       *defn = strdup(te->defn);
2214         char       *ptr1;
2215         char       *ptr2 = NULL;
2216         int                     encoding;
2217
2218         ptr1 = strchr(defn, '\'');
2219         if (ptr1)
2220                 ptr2 = strchr(++ptr1, '\'');
2221         if (ptr2)
2222         {
2223                 *ptr2 = '\0';
2224                 encoding = pg_char_to_encoding(ptr1);
2225                 if (encoding < 0)
2226                         die_horribly(AH, modulename, "unrecognized encoding \"%s\"\n",
2227                                                  ptr1);
2228                 AH->public.encoding = encoding;
2229         }
2230         else
2231                 die_horribly(AH, modulename, "invalid ENCODING item: %s\n",
2232                                          te->defn);
2233
2234         free(defn);
2235 }
2236
2237 static void
2238 processStdStringsEntry(ArchiveHandle *AH, TocEntry *te)
2239 {
2240         /* te->defn should have the form SET standard_conforming_strings = 'x'; */
2241         char       *ptr1;
2242
2243         ptr1 = strchr(te->defn, '\'');
2244         if (ptr1 && strncmp(ptr1, "'on'", 4) == 0)
2245                 AH->public.std_strings = true;
2246         else if (ptr1 && strncmp(ptr1, "'off'", 5) == 0)
2247                 AH->public.std_strings = false;
2248         else
2249                 die_horribly(AH, modulename, "invalid STDSTRINGS item: %s\n",
2250                                          te->defn);
2251 }
2252
2253 static teReqs
2254 _tocEntryRequired(TocEntry *te, RestoreOptions *ropt, bool include_acls)
2255 {
2256         teReqs          res = REQ_ALL;
2257
2258         /* ENCODING and STDSTRINGS items are dumped specially, so always reject */
2259         if (strcmp(te->desc, "ENCODING") == 0 ||
2260                 strcmp(te->desc, "STDSTRINGS") == 0)
2261                 return 0;
2262
2263         /* If it's an ACL, maybe ignore it */
2264         if ((!include_acls || ropt->aclsSkip) && _tocEntryIsACL(te))
2265                 return 0;
2266
2267         /* Ignore DATABASE entry unless we should create it */
2268         if (!ropt->createDB && strcmp(te->desc, "DATABASE") == 0)
2269                 return 0;
2270
2271         /* Check options for selective dump/restore */
2272         if (ropt->schemaNames)
2273         {
2274                 /* If no namespace is specified, it means all. */
2275                 if (!te->namespace)
2276                         return 0;
2277                 if (strcmp(ropt->schemaNames, te->namespace) != 0)
2278                         return 0;
2279         }
2280
2281         if (ropt->selTypes)
2282         {
2283                 if (strcmp(te->desc, "TABLE") == 0 ||
2284                         strcmp(te->desc, "TABLE DATA") == 0)
2285                 {
2286                         if (!ropt->selTable)
2287                                 return 0;
2288                         if (ropt->tableNames && strcmp(ropt->tableNames, te->tag) != 0)
2289                                 return 0;
2290                 }
2291                 else if (strcmp(te->desc, "INDEX") == 0)
2292                 {
2293                         if (!ropt->selIndex)
2294                                 return 0;
2295                         if (ropt->indexNames && strcmp(ropt->indexNames, te->tag) != 0)
2296                                 return 0;
2297                 }
2298                 else if (strcmp(te->desc, "FUNCTION") == 0)
2299                 {
2300                         if (!ropt->selFunction)
2301                                 return 0;
2302                         if (ropt->functionNames && strcmp(ropt->functionNames, te->tag) != 0)
2303                                 return 0;
2304                 }
2305                 else if (strcmp(te->desc, "TRIGGER") == 0)
2306                 {
2307                         if (!ropt->selTrigger)
2308                                 return 0;
2309                         if (ropt->triggerNames && strcmp(ropt->triggerNames, te->tag) != 0)
2310                                 return 0;
2311                 }
2312                 else
2313                         return 0;
2314         }
2315
2316         /*
2317          * Check if we had a dataDumper. Indicates if the entry is schema or data
2318          */
2319         if (!te->hadDumper)
2320         {
2321                 /*
2322                  * Special Case: If 'SEQUENCE SET' or anything to do with BLOBs, then
2323                  * it is considered a data entry.  We don't need to check for the
2324                  * BLOBS entry or old-style BLOB COMMENTS, because they will have
2325                  * hadDumper = true ... but we do need to check new-style BLOB
2326                  * comments.
2327                  */
2328                 if (strcmp(te->desc, "SEQUENCE SET") == 0 ||
2329                         strcmp(te->desc, "BLOB") == 0 ||
2330                         (strcmp(te->desc, "ACL") == 0 &&
2331                          strncmp(te->tag, "LARGE OBJECT ", 13) == 0) ||
2332                         (strcmp(te->desc, "COMMENT") == 0 &&
2333                          strncmp(te->tag, "LARGE OBJECT ", 13) == 0))
2334                         res = res & REQ_DATA;
2335                 else
2336                         res = res & ~REQ_DATA;
2337         }
2338
2339         /*
2340          * Special case: <Init> type with <Max OID> tag; this is obsolete and we
2341          * always ignore it.
2342          */
2343         if ((strcmp(te->desc, "<Init>") == 0) && (strcmp(te->tag, "Max OID") == 0))
2344                 return 0;
2345
2346         /* Mask it if we only want schema */
2347         if (ropt->schemaOnly)
2348                 res = res & REQ_SCHEMA;
2349
2350         /* Mask it we only want data */
2351         if (ropt->dataOnly)
2352                 res = res & REQ_DATA;
2353
2354         /* Mask it if we don't have a schema contribution */
2355         if (!te->defn || strlen(te->defn) == 0)
2356                 res = res & ~REQ_SCHEMA;
2357
2358         /* Finally, if there's a per-ID filter, limit based on that as well */
2359         if (ropt->idWanted && !ropt->idWanted[te->dumpId - 1])
2360                 return 0;
2361
2362         return res;
2363 }
2364
2365 /*
2366  * Identify TOC entries that are ACLs.
2367  */
2368 static bool
2369 _tocEntryIsACL(TocEntry *te)
2370 {
2371         /* "ACL LANGUAGE" was a crock emitted only in PG 7.4 */
2372         if (strcmp(te->desc, "ACL") == 0 ||
2373                 strcmp(te->desc, "ACL LANGUAGE") == 0 ||
2374                 strcmp(te->desc, "DEFAULT ACL") == 0)
2375                 return true;
2376         return false;
2377 }
2378
2379 /*
2380  * Issue SET commands for parameters that we want to have set the same way
2381  * at all times during execution of a restore script.
2382  */
2383 static void
2384 _doSetFixedOutputState(ArchiveHandle *AH)
2385 {
2386         /* Disable statement_timeout in archive for pg_restore/psql  */
2387         ahprintf(AH, "SET statement_timeout = 0;\n");
2388
2389         /* Select the correct character set encoding */
2390         ahprintf(AH, "SET client_encoding = '%s';\n",
2391                          pg_encoding_to_char(AH->public.encoding));
2392
2393         /* Select the correct string literal syntax */
2394         ahprintf(AH, "SET standard_conforming_strings = %s;\n",
2395                          AH->public.std_strings ? "on" : "off");
2396
2397         /* Select the role to be used during restore */
2398         if (AH->ropt && AH->ropt->use_role)
2399                 ahprintf(AH, "SET ROLE %s;\n", fmtId(AH->ropt->use_role));
2400
2401         /* Make sure function checking is disabled */
2402         ahprintf(AH, "SET check_function_bodies = false;\n");
2403
2404         /* Avoid annoying notices etc */
2405         ahprintf(AH, "SET client_min_messages = warning;\n");
2406         if (!AH->public.std_strings)
2407                 ahprintf(AH, "SET escape_string_warning = off;\n");
2408
2409         ahprintf(AH, "\n");
2410 }
2411
2412 /*
2413  * Issue a SET SESSION AUTHORIZATION command.  Caller is responsible
2414  * for updating state if appropriate.  If user is NULL or an empty string,
2415  * the specification DEFAULT will be used.
2416  */
2417 static void
2418 _doSetSessionAuth(ArchiveHandle *AH, const char *user)
2419 {
2420         PQExpBuffer cmd = createPQExpBuffer();
2421
2422         appendPQExpBuffer(cmd, "SET SESSION AUTHORIZATION ");
2423
2424         /*
2425          * SQL requires a string literal here.  Might as well be correct.
2426          */
2427         if (user && *user)
2428                 appendStringLiteralAHX(cmd, user, AH);
2429         else
2430                 appendPQExpBuffer(cmd, "DEFAULT");
2431         appendPQExpBuffer(cmd, ";");
2432
2433         if (RestoringToDB(AH))
2434         {
2435                 PGresult   *res;
2436
2437                 res = PQexec(AH->connection, cmd->data);
2438
2439                 if (!res || PQresultStatus(res) != PGRES_COMMAND_OK)
2440                         /* NOT warn_or_die_horribly... use -O instead to skip this. */
2441                         die_horribly(AH, modulename, "could not set session user to \"%s\": %s",
2442                                                  user, PQerrorMessage(AH->connection));
2443
2444                 PQclear(res);
2445         }
2446         else
2447                 ahprintf(AH, "%s\n\n", cmd->data);
2448
2449         destroyPQExpBuffer(cmd);
2450 }
2451
2452
2453 /*
2454  * Issue a SET default_with_oids command.  Caller is responsible
2455  * for updating state if appropriate.
2456  */
2457 static void
2458 _doSetWithOids(ArchiveHandle *AH, const bool withOids)
2459 {
2460         PQExpBuffer cmd = createPQExpBuffer();
2461
2462         appendPQExpBuffer(cmd, "SET default_with_oids = %s;", withOids ?
2463                                           "true" : "false");
2464
2465         if (RestoringToDB(AH))
2466         {
2467                 PGresult   *res;
2468
2469                 res = PQexec(AH->connection, cmd->data);
2470
2471                 if (!res || PQresultStatus(res) != PGRES_COMMAND_OK)
2472                         warn_or_die_horribly(AH, modulename,
2473                                                                  "could not set default_with_oids: %s",
2474                                                                  PQerrorMessage(AH->connection));
2475
2476                 PQclear(res);
2477         }
2478         else
2479                 ahprintf(AH, "%s\n\n", cmd->data);
2480
2481         destroyPQExpBuffer(cmd);
2482 }
2483
2484
2485 /*
2486  * Issue the commands to connect to the specified database.
2487  *
2488  * If we're currently restoring right into a database, this will
2489  * actually establish a connection. Otherwise it puts a \connect into
2490  * the script output.
2491  *
2492  * NULL dbname implies reconnecting to the current DB (pretty useless).
2493  */
2494 static void
2495 _reconnectToDB(ArchiveHandle *AH, const char *dbname)
2496 {
2497         if (RestoringToDB(AH))
2498                 ReconnectToServer(AH, dbname, NULL);
2499         else
2500         {
2501                 PQExpBuffer qry = createPQExpBuffer();
2502
2503                 appendPQExpBuffer(qry, "\\connect %s\n\n",
2504                                                   dbname ? fmtId(dbname) : "-");
2505                 ahprintf(AH, "%s", qry->data);
2506                 destroyPQExpBuffer(qry);
2507         }
2508
2509         /*
2510          * NOTE: currUser keeps track of what the imaginary session user in our
2511          * script is.  It's now effectively reset to the original userID.
2512          */
2513         if (AH->currUser)
2514                 free(AH->currUser);
2515         AH->currUser = NULL;
2516
2517         /* don't assume we still know the output schema, tablespace, etc either */
2518         if (AH->currSchema)
2519                 free(AH->currSchema);
2520         AH->currSchema = NULL;
2521         if (AH->currTablespace)
2522                 free(AH->currTablespace);
2523         AH->currTablespace = NULL;
2524         AH->currWithOids = -1;
2525
2526         /* re-establish fixed state */
2527         _doSetFixedOutputState(AH);
2528 }
2529
2530 /*
2531  * Become the specified user, and update state to avoid redundant commands
2532  *
2533  * NULL or empty argument is taken to mean restoring the session default
2534  */
2535 static void
2536 _becomeUser(ArchiveHandle *AH, const char *user)
2537 {
2538         if (!user)
2539                 user = "";                              /* avoid null pointers */
2540
2541         if (AH->currUser && strcmp(AH->currUser, user) == 0)
2542                 return;                                 /* no need to do anything */
2543
2544         _doSetSessionAuth(AH, user);
2545
2546         /*
2547          * NOTE: currUser keeps track of what the imaginary session user in our
2548          * script is
2549          */
2550         if (AH->currUser)
2551                 free(AH->currUser);
2552         AH->currUser = strdup(user);
2553 }
2554
2555 /*
2556  * Become the owner of the given TOC entry object.      If
2557  * changes in ownership are not allowed, this doesn't do anything.
2558  */
2559 static void
2560 _becomeOwner(ArchiveHandle *AH, TocEntry *te)
2561 {
2562         if (AH->ropt && (AH->ropt->noOwner || !AH->ropt->use_setsessauth))
2563                 return;
2564
2565         _becomeUser(AH, te->owner);
2566 }
2567
2568
2569 /*
2570  * Set the proper default_with_oids value for the table.
2571  */
2572 static void
2573 _setWithOids(ArchiveHandle *AH, TocEntry *te)
2574 {
2575         if (AH->currWithOids != te->withOids)
2576         {
2577                 _doSetWithOids(AH, te->withOids);
2578                 AH->currWithOids = te->withOids;
2579         }
2580 }
2581
2582
2583 /*
2584  * Issue the commands to select the specified schema as the current schema
2585  * in the target database.
2586  */
2587 static void
2588 _selectOutputSchema(ArchiveHandle *AH, const char *schemaName)
2589 {
2590         PQExpBuffer qry;
2591
2592         if (!schemaName || *schemaName == '\0' ||
2593                 (AH->currSchema && strcmp(AH->currSchema, schemaName) == 0))
2594                 return;                                 /* no need to do anything */
2595
2596         qry = createPQExpBuffer();
2597
2598         appendPQExpBuffer(qry, "SET search_path = %s",
2599                                           fmtId(schemaName));
2600         if (strcmp(schemaName, "pg_catalog") != 0)
2601                 appendPQExpBuffer(qry, ", pg_catalog");
2602
2603         if (RestoringToDB(AH))
2604         {
2605                 PGresult   *res;
2606
2607                 res = PQexec(AH->connection, qry->data);
2608
2609                 if (!res || PQresultStatus(res) != PGRES_COMMAND_OK)
2610                         warn_or_die_horribly(AH, modulename,
2611                                                                  "could not set search_path to \"%s\": %s",
2612                                                                  schemaName, PQerrorMessage(AH->connection));
2613
2614                 PQclear(res);
2615         }
2616         else
2617                 ahprintf(AH, "%s;\n\n", qry->data);
2618
2619         if (AH->currSchema)
2620                 free(AH->currSchema);
2621         AH->currSchema = strdup(schemaName);
2622
2623         destroyPQExpBuffer(qry);
2624 }
2625
2626 /*
2627  * Issue the commands to select the specified tablespace as the current one
2628  * in the target database.
2629  */
2630 static void
2631 _selectTablespace(ArchiveHandle *AH, const char *tablespace)
2632 {
2633         PQExpBuffer qry;
2634         const char *want,
2635                            *have;
2636
2637         /* do nothing in --no-tablespaces mode */
2638         if (AH->ropt->noTablespace)
2639                 return;
2640
2641         have = AH->currTablespace;
2642         want = tablespace;
2643
2644         /* no need to do anything for non-tablespace object */
2645         if (!want)
2646                 return;
2647
2648         if (have && strcmp(want, have) == 0)
2649                 return;                                 /* no need to do anything */
2650
2651         qry = createPQExpBuffer();
2652
2653         if (strcmp(want, "") == 0)
2654         {
2655                 /* We want the tablespace to be the database's default */
2656                 appendPQExpBuffer(qry, "SET default_tablespace = ''");
2657         }
2658         else
2659         {
2660                 /* We want an explicit tablespace */
2661                 appendPQExpBuffer(qry, "SET default_tablespace = %s", fmtId(want));
2662         }
2663
2664         if (RestoringToDB(AH))
2665         {
2666                 PGresult   *res;
2667
2668                 res = PQexec(AH->connection, qry->data);
2669
2670                 if (!res || PQresultStatus(res) != PGRES_COMMAND_OK)
2671                         warn_or_die_horribly(AH, modulename,
2672                                                                  "could not set default_tablespace to %s: %s",
2673                                                                  fmtId(want), PQerrorMessage(AH->connection));
2674
2675                 PQclear(res);
2676         }
2677         else
2678                 ahprintf(AH, "%s;\n\n", qry->data);
2679
2680         if (AH->currTablespace)
2681                 free(AH->currTablespace);
2682         AH->currTablespace = strdup(want);
2683
2684         destroyPQExpBuffer(qry);
2685 }
2686
2687 /*
2688  * Extract an object description for a TOC entry, and append it to buf.
2689  *
2690  * This is not quite as general as it may seem, since it really only
2691  * handles constructing the right thing to put into ALTER ... OWNER TO.
2692  *
2693  * The whole thing is pretty grotty, but we are kind of stuck since the
2694  * information used is all that's available in older dump files.
2695  */
2696 static void
2697 _getObjectDescription(PQExpBuffer buf, TocEntry *te, ArchiveHandle *AH)
2698 {
2699         const char *type = te->desc;
2700
2701         /* Use ALTER TABLE for views and sequences */
2702         if (strcmp(type, "VIEW") == 0 || strcmp(type, "SEQUENCE") == 0)
2703                 type = "TABLE";
2704
2705         /* objects named by a schema and name */
2706         if (strcmp(type, "CONVERSION") == 0 ||
2707                 strcmp(type, "DOMAIN") == 0 ||
2708                 strcmp(type, "TABLE") == 0 ||
2709                 strcmp(type, "TYPE") == 0 ||
2710                 strcmp(type, "TEXT SEARCH DICTIONARY") == 0 ||
2711                 strcmp(type, "TEXT SEARCH CONFIGURATION") == 0)
2712         {
2713                 appendPQExpBuffer(buf, "%s ", type);
2714                 if (te->namespace && te->namespace[0])  /* is null pre-7.3 */
2715                         appendPQExpBuffer(buf, "%s.", fmtId(te->namespace));
2716
2717                 /*
2718                  * Pre-7.3 pg_dump would sometimes (not always) put a fmtId'd name
2719                  * into te->tag for an index. This check is heuristic, so make its
2720                  * scope as narrow as possible.
2721                  */
2722                 if (AH->version < K_VERS_1_7 &&
2723                         te->tag[0] == '"' &&
2724                         te->tag[strlen(te->tag) - 1] == '"' &&
2725                         strcmp(type, "INDEX") == 0)
2726                         appendPQExpBuffer(buf, "%s", te->tag);
2727                 else
2728                         appendPQExpBuffer(buf, "%s", fmtId(te->tag));
2729                 return;
2730         }
2731
2732         /* objects named by just a name */
2733         if (strcmp(type, "DATABASE") == 0 ||
2734                 strcmp(type, "PROCEDURAL LANGUAGE") == 0 ||
2735                 strcmp(type, "SCHEMA") == 0 ||
2736                 strcmp(type, "FOREIGN DATA WRAPPER") == 0 ||
2737                 strcmp(type, "SERVER") == 0 ||
2738                 strcmp(type, "USER MAPPING") == 0)
2739         {
2740                 appendPQExpBuffer(buf, "%s %s", type, fmtId(te->tag));
2741                 return;
2742         }
2743
2744         /* BLOBs just have a name, but it's numeric so must not use fmtId */
2745         if (strcmp(type, "BLOB") == 0)
2746         {
2747                 appendPQExpBuffer(buf, "LARGE OBJECT %s", te->tag);
2748                 return;
2749         }
2750
2751         /*
2752          * These object types require additional decoration.  Fortunately, the
2753          * information needed is exactly what's in the DROP command.
2754          */
2755         if (strcmp(type, "AGGREGATE") == 0 ||
2756                 strcmp(type, "FUNCTION") == 0 ||
2757                 strcmp(type, "OPERATOR") == 0 ||
2758                 strcmp(type, "OPERATOR CLASS") == 0 ||
2759                 strcmp(type, "OPERATOR FAMILY") == 0)
2760         {
2761                 /* Chop "DROP " off the front and make a modifiable copy */
2762                 char       *first = strdup(te->dropStmt + 5);
2763                 char       *last;
2764
2765                 /* point to last character in string */
2766                 last = first + strlen(first) - 1;
2767
2768                 /* Strip off any ';' or '\n' at the end */
2769                 while (last >= first && (*last == '\n' || *last == ';'))
2770                         last--;
2771                 *(last + 1) = '\0';
2772
2773                 appendPQExpBufferStr(buf, first);
2774
2775                 free(first);
2776                 return;
2777         }
2778
2779         write_msg(modulename, "WARNING: don't know how to set owner for object type %s\n",
2780                           type);
2781 }
2782
2783 static void
2784 _printTocEntry(ArchiveHandle *AH, TocEntry *te, RestoreOptions *ropt, bool isData, bool acl_pass)
2785 {
2786         /* ACLs are dumped only during acl pass */
2787         if (acl_pass)
2788         {
2789                 if (!_tocEntryIsACL(te))
2790                         return;
2791         }
2792         else
2793         {
2794                 if (_tocEntryIsACL(te))
2795                         return;
2796         }
2797
2798         /*
2799          * Avoid dumping the public schema, as it will already be created ...
2800          * unless we are using --clean mode, in which case it's been deleted and
2801          * we'd better recreate it.  Likewise for its comment, if any.
2802          */
2803         if (!ropt->dropSchema)
2804         {
2805                 if (strcmp(te->desc, "SCHEMA") == 0 &&
2806                         strcmp(te->tag, "public") == 0)
2807                         return;
2808                 /* The comment restore would require super-user privs, so avoid it. */
2809                 if (strcmp(te->desc, "COMMENT") == 0 &&
2810                         strcmp(te->tag, "SCHEMA public") == 0)
2811                         return;
2812         }
2813
2814         /* Select owner, schema, and tablespace as necessary */
2815         _becomeOwner(AH, te);
2816         _selectOutputSchema(AH, te->namespace);
2817         _selectTablespace(AH, te->tablespace);
2818
2819         /* Set up OID mode too */
2820         if (strcmp(te->desc, "TABLE") == 0)
2821                 _setWithOids(AH, te);
2822
2823         /* Emit header comment for item */
2824         if (!AH->noTocComments)
2825         {
2826                 const char *pfx;
2827
2828                 if (isData)
2829                         pfx = "Data for ";
2830                 else
2831                         pfx = "";
2832
2833                 ahprintf(AH, "--\n");
2834                 if (AH->public.verbose)
2835                 {
2836                         ahprintf(AH, "-- TOC entry %d (class %u OID %u)\n",
2837                                          te->dumpId, te->catalogId.tableoid, te->catalogId.oid);
2838                         if (te->nDeps > 0)
2839                         {
2840                                 int                     i;
2841
2842                                 ahprintf(AH, "-- Dependencies:");
2843                                 for (i = 0; i < te->nDeps; i++)
2844                                         ahprintf(AH, " %d", te->dependencies[i]);
2845                                 ahprintf(AH, "\n");
2846                         }
2847                 }
2848                 ahprintf(AH, "-- %sName: %s; Type: %s; Schema: %s; Owner: %s",
2849                                  pfx, te->tag, te->desc,
2850                                  te->namespace ? te->namespace : "-",
2851                                  ropt->noOwner ? "-" : te->owner);
2852                 if (te->tablespace && !ropt->noTablespace)
2853                         ahprintf(AH, "; Tablespace: %s", te->tablespace);
2854                 ahprintf(AH, "\n");
2855
2856                 if (AH->PrintExtraTocPtr !=NULL)
2857                         (*AH->PrintExtraTocPtr) (AH, te);
2858                 ahprintf(AH, "--\n\n");
2859         }
2860
2861         /*
2862          * Actually print the definition.
2863          *
2864          * Really crude hack for suppressing AUTHORIZATION clause that old pg_dump
2865          * versions put into CREATE SCHEMA.  We have to do this when --no-owner
2866          * mode is selected.  This is ugly, but I see no other good way ...
2867          */
2868         if (ropt->noOwner && strcmp(te->desc, "SCHEMA") == 0)
2869         {
2870                 ahprintf(AH, "CREATE SCHEMA %s;\n\n\n", fmtId(te->tag));
2871         }
2872         else
2873         {
2874                 if (strlen(te->defn) > 0)
2875                         ahprintf(AH, "%s\n\n", te->defn);
2876         }
2877
2878         /*
2879          * If we aren't using SET SESSION AUTH to determine ownership, we must
2880          * instead issue an ALTER OWNER command.  We assume that anything without
2881          * a DROP command is not a separately ownable object.  All the categories
2882          * with DROP commands must appear in one list or the other.
2883          */
2884         if (!ropt->noOwner && !ropt->use_setsessauth &&
2885                 strlen(te->owner) > 0 && strlen(te->dropStmt) > 0)
2886         {
2887                 if (strcmp(te->desc, "AGGREGATE") == 0 ||
2888                         strcmp(te->desc, "BLOB") == 0 ||
2889                         strcmp(te->desc, "CONVERSION") == 0 ||
2890                         strcmp(te->desc, "DATABASE") == 0 ||
2891                         strcmp(te->desc, "DOMAIN") == 0 ||
2892                         strcmp(te->desc, "FUNCTION") == 0 ||
2893                         strcmp(te->desc, "OPERATOR") == 0 ||
2894                         strcmp(te->desc, "OPERATOR CLASS") == 0 ||
2895                         strcmp(te->desc, "OPERATOR FAMILY") == 0 ||
2896                         strcmp(te->desc, "PROCEDURAL LANGUAGE") == 0 ||
2897                         strcmp(te->desc, "SCHEMA") == 0 ||
2898                         strcmp(te->desc, "TABLE") == 0 ||
2899                         strcmp(te->desc, "TYPE") == 0 ||
2900                         strcmp(te->desc, "VIEW") == 0 ||
2901                         strcmp(te->desc, "SEQUENCE") == 0 ||
2902                         strcmp(te->desc, "TEXT SEARCH DICTIONARY") == 0 ||
2903                         strcmp(te->desc, "TEXT SEARCH CONFIGURATION") == 0 ||
2904                         strcmp(te->desc, "FOREIGN DATA WRAPPER") == 0 ||
2905                         strcmp(te->desc, "SERVER") == 0)
2906                 {
2907                         PQExpBuffer temp = createPQExpBuffer();
2908
2909                         appendPQExpBuffer(temp, "ALTER ");
2910                         _getObjectDescription(temp, te, AH);
2911                         appendPQExpBuffer(temp, " OWNER TO %s;", fmtId(te->owner));
2912                         ahprintf(AH, "%s\n\n", temp->data);
2913                         destroyPQExpBuffer(temp);
2914                 }
2915                 else if (strcmp(te->desc, "CAST") == 0 ||
2916                                  strcmp(te->desc, "CHECK CONSTRAINT") == 0 ||
2917                                  strcmp(te->desc, "CONSTRAINT") == 0 ||
2918                                  strcmp(te->desc, "DEFAULT") == 0 ||
2919                                  strcmp(te->desc, "FK CONSTRAINT") == 0 ||
2920                                  strcmp(te->desc, "INDEX") == 0 ||
2921                                  strcmp(te->desc, "RULE") == 0 ||
2922                                  strcmp(te->desc, "TRIGGER") == 0 ||
2923                                  strcmp(te->desc, "USER MAPPING") == 0)
2924                 {
2925                         /* these object types don't have separate owners */
2926                 }
2927                 else
2928                 {
2929                         write_msg(modulename, "WARNING: don't know how to set owner for object type %s\n",
2930                                           te->desc);
2931                 }
2932         }
2933
2934         /*
2935          * If it's an ACL entry, it might contain SET SESSION AUTHORIZATION
2936          * commands, so we can no longer assume we know the current auth setting.
2937          */
2938         if (acl_pass)
2939         {
2940                 if (AH->currUser)
2941                         free(AH->currUser);
2942                 AH->currUser = NULL;
2943         }
2944 }
2945
2946 void
2947 WriteHead(ArchiveHandle *AH)
2948 {
2949         struct tm       crtm;
2950
2951         (*AH->WriteBufPtr) (AH, "PGDMP", 5);            /* Magic code */
2952         (*AH->WriteBytePtr) (AH, AH->vmaj);
2953         (*AH->WriteBytePtr) (AH, AH->vmin);
2954         (*AH->WriteBytePtr) (AH, AH->vrev);
2955         (*AH->WriteBytePtr) (AH, AH->intSize);
2956         (*AH->WriteBytePtr) (AH, AH->offSize);
2957         (*AH->WriteBytePtr) (AH, AH->format);
2958
2959 #ifndef HAVE_LIBZ
2960         if (AH->compression != 0)
2961                 write_msg(modulename, "WARNING: requested compression not available in this "
2962                                   "installation -- archive will be uncompressed\n");
2963
2964         AH->compression = 0;
2965 #endif
2966
2967         WriteInt(AH, AH->compression);
2968
2969         crtm = *localtime(&AH->createDate);
2970         WriteInt(AH, crtm.tm_sec);
2971         WriteInt(AH, crtm.tm_min);
2972         WriteInt(AH, crtm.tm_hour);
2973         WriteInt(AH, crtm.tm_mday);
2974         WriteInt(AH, crtm.tm_mon);
2975         WriteInt(AH, crtm.tm_year);
2976         WriteInt(AH, crtm.tm_isdst);
2977         WriteStr(AH, PQdb(AH->connection));
2978         WriteStr(AH, AH->public.remoteVersionStr);
2979         WriteStr(AH, PG_VERSION);
2980 }
2981
2982 void
2983 ReadHead(ArchiveHandle *AH)
2984 {
2985         char            tmpMag[7];
2986         int                     fmt;
2987         struct tm       crtm;
2988
2989         /*
2990          * If we haven't already read the header, do so.
2991          *
2992          * NB: this code must agree with _discoverArchiveFormat().      Maybe find a
2993          * way to unify the cases?
2994          */
2995         if (!AH->readHeader)
2996         {
2997                 if ((*AH->ReadBufPtr) (AH, tmpMag, 5) != 5)
2998                         die_horribly(AH, modulename, "unexpected end of file\n");
2999
3000                 if (strncmp(tmpMag, "PGDMP", 5) != 0)
3001                         die_horribly(AH, modulename, "did not find magic string in file header\n");
3002
3003                 AH->vmaj = (*AH->ReadBytePtr) (AH);
3004                 AH->vmin = (*AH->ReadBytePtr) (AH);
3005
3006                 if (AH->vmaj > 1 || ((AH->vmaj == 1) && (AH->vmin > 0)))                /* Version > 1.0 */
3007                         AH->vrev = (*AH->ReadBytePtr) (AH);
3008                 else
3009                         AH->vrev = 0;
3010
3011                 AH->version = ((AH->vmaj * 256 + AH->vmin) * 256 + AH->vrev) * 256 + 0;
3012
3013                 if (AH->version < K_VERS_1_0 || AH->version > K_VERS_MAX)
3014                         die_horribly(AH, modulename, "unsupported version (%d.%d) in file header\n",
3015                                                  AH->vmaj, AH->vmin);
3016
3017                 AH->intSize = (*AH->ReadBytePtr) (AH);
3018                 if (AH->intSize > 32)
3019                         die_horribly(AH, modulename, "sanity check on integer size (%lu) failed\n",
3020                                                  (unsigned long) AH->intSize);
3021
3022                 if (AH->intSize > sizeof(int))
3023                         write_msg(modulename, "WARNING: archive was made on a machine with larger integers, some operations might fail\n");
3024
3025                 if (AH->version >= K_VERS_1_7)
3026                         AH->offSize = (*AH->ReadBytePtr) (AH);
3027                 else
3028                         AH->offSize = AH->intSize;
3029
3030                 fmt = (*AH->ReadBytePtr) (AH);
3031
3032                 if (AH->format != fmt)
3033                         die_horribly(AH, modulename, "expected format (%d) differs from format found in file (%d)\n",
3034                                                  AH->format, fmt);
3035         }
3036
3037         if (AH->version >= K_VERS_1_2)
3038         {
3039                 if (AH->version < K_VERS_1_4)
3040                         AH->compression = (*AH->ReadBytePtr) (AH);
3041                 else
3042                         AH->compression = ReadInt(AH);
3043         }
3044         else
3045                 AH->compression = Z_DEFAULT_COMPRESSION;
3046
3047 #ifndef HAVE_LIBZ
3048         if (AH->compression != 0)
3049                 write_msg(modulename, "WARNING: archive is compressed, but this installation does not support compression -- no data will be available\n");
3050 #endif
3051
3052         if (AH->version >= K_VERS_1_4)
3053         {
3054                 crtm.tm_sec = ReadInt(AH);
3055                 crtm.tm_min = ReadInt(AH);
3056                 crtm.tm_hour = ReadInt(AH);
3057                 crtm.tm_mday = ReadInt(AH);
3058                 crtm.tm_mon = ReadInt(AH);
3059                 crtm.tm_year = ReadInt(AH);
3060                 crtm.tm_isdst = ReadInt(AH);
3061
3062                 AH->archdbname = ReadStr(AH);
3063
3064                 AH->createDate = mktime(&crtm);
3065
3066                 if (AH->createDate == (time_t) -1)
3067                         write_msg(modulename, "WARNING: invalid creation date in header\n");
3068         }
3069
3070         if (AH->version >= K_VERS_1_10)
3071         {
3072                 AH->archiveRemoteVersion = ReadStr(AH);
3073                 AH->archiveDumpVersion = ReadStr(AH);
3074         }
3075 }
3076
3077
3078 /*
3079  * checkSeek
3080  *        check to see if ftell/fseek can be performed.
3081  */
3082 bool
3083 checkSeek(FILE *fp)
3084 {
3085         pgoff_t         tpos;
3086
3087         /*
3088          * If pgoff_t is wider than long, we must have "real" fseeko and not an
3089          * emulation using fseek.  Otherwise report no seek capability.
3090          */
3091 #ifndef HAVE_FSEEKO
3092         if (sizeof(pgoff_t) > sizeof(long))
3093                 return false;
3094 #endif
3095
3096         /* Check that ftello works on this file */
3097         errno = 0;
3098         tpos = ftello(fp);
3099         if (errno)
3100                 return false;
3101
3102         /*
3103          * Check that fseeko(SEEK_SET) works, too.      NB: we used to try to test
3104          * this with fseeko(fp, 0, SEEK_CUR).  But some platforms treat that as a
3105          * successful no-op even on files that are otherwise unseekable.
3106          */
3107         if (fseeko(fp, tpos, SEEK_SET) != 0)
3108                 return false;
3109
3110         return true;
3111 }
3112
3113
3114 /*
3115  * dumpTimestamp
3116  */
3117 static void
3118 dumpTimestamp(ArchiveHandle *AH, const char *msg, time_t tim)
3119 {
3120         char            buf[256];
3121
3122         /*
3123          * We don't print the timezone on Win32, because the names are long and
3124          * localized, which means they may contain characters in various random
3125          * encodings; this has been seen to cause encoding errors when reading the
3126          * dump script.
3127          */
3128         if (strftime(buf, sizeof(buf),
3129 #ifndef WIN32
3130                                  "%Y-%m-%d %H:%M:%S %Z",
3131 #else
3132                                  "%Y-%m-%d %H:%M:%S",
3133 #endif
3134                                  localtime(&tim)) != 0)
3135                 ahprintf(AH, "-- %s %s\n\n", msg, buf);
3136 }
3137
3138
3139 /*
3140  * Main engine for parallel restore.
3141  *
3142  * Work is done in three phases.
3143  * First we process tocEntries until we come to one that is marked
3144  * SECTION_DATA or SECTION_POST_DATA, in a single connection, just as for a
3145  * standard restore.  Second we process the remaining non-ACL steps in
3146  * parallel worker children (threads on Windows, processes on Unix), each of
3147  * which connects separately to the database.  Finally we process all the ACL
3148  * entries in a single connection (that happens back in RestoreArchive).
3149  */
3150 static void
3151 restore_toc_entries_parallel(ArchiveHandle *AH)
3152 {
3153         RestoreOptions *ropt = AH->ropt;
3154         int                     n_slots = ropt->number_of_jobs;
3155         ParallelSlot *slots;
3156         int                     work_status;
3157         int                     next_slot;
3158         TocEntry        pending_list;
3159         TocEntry        ready_list;
3160         TocEntry   *next_work_item;
3161         thandle         ret_child;
3162         TocEntry   *te;
3163
3164         ahlog(AH, 2, "entering restore_toc_entries_parallel\n");
3165
3166         /* we haven't got round to making this work for all archive formats */
3167         if (AH->ClonePtr == NULL || AH->ReopenPtr == NULL)
3168                 die_horribly(AH, modulename, "parallel restore is not supported with this archive file format\n");
3169
3170         /* doesn't work if the archive represents dependencies as OIDs, either */
3171         if (AH->version < K_VERS_1_8)
3172                 die_horribly(AH, modulename, "parallel restore is not supported with archives made by pre-8.0 pg_dump\n");
3173
3174         slots = (ParallelSlot *) calloc(sizeof(ParallelSlot), n_slots);
3175
3176         /* Adjust dependency information */
3177         fix_dependencies(AH);
3178
3179         /*
3180          * Do all the early stuff in a single connection in the parent. There's no
3181          * great point in running it in parallel, in fact it will actually run
3182          * faster in a single connection because we avoid all the connection and
3183          * setup overhead.
3184          */
3185         for (next_work_item = AH->toc->next; next_work_item != AH->toc; next_work_item = next_work_item->next)
3186         {
3187                 if (next_work_item->section == SECTION_DATA ||
3188                         next_work_item->section == SECTION_POST_DATA)
3189                         break;
3190
3191                 ahlog(AH, 1, "processing item %d %s %s\n",
3192                           next_work_item->dumpId,
3193                           next_work_item->desc, next_work_item->tag);
3194
3195                 (void) restore_toc_entry(AH, next_work_item, ropt, false);
3196
3197                 /* there should be no touch of ready_list here, so pass NULL */
3198                 reduce_dependencies(AH, next_work_item, NULL);
3199         }
3200
3201         /*
3202          * Now close parent connection in prep for parallel steps.      We do this
3203          * mainly to ensure that we don't exceed the specified number of parallel
3204          * connections.
3205          */
3206         PQfinish(AH->connection);
3207         AH->connection = NULL;
3208
3209         /* blow away any transient state from the old connection */
3210         if (AH->currUser)
3211                 free(AH->currUser);
3212         AH->currUser = NULL;
3213         if (AH->currSchema)
3214                 free(AH->currSchema);
3215         AH->currSchema = NULL;
3216         if (AH->currTablespace)
3217                 free(AH->currTablespace);
3218         AH->currTablespace = NULL;
3219         AH->currWithOids = -1;
3220
3221         /*
3222          * Initialize the lists of pending and ready items.  After this setup, the
3223          * pending list is everything that needs to be done but is blocked by one
3224          * or more dependencies, while the ready list contains items that have no
3225          * remaining dependencies.      Note: we don't yet filter out entries that
3226          * aren't going to be restored.  They might participate in dependency
3227          * chains connecting entries that should be restored, so we treat them as
3228          * live until we actually process them.
3229          */
3230         par_list_header_init(&pending_list);
3231         par_list_header_init(&ready_list);
3232         for (; next_work_item != AH->toc; next_work_item = next_work_item->next)
3233         {
3234                 if (next_work_item->depCount > 0)
3235                         par_list_append(&pending_list, next_work_item);
3236                 else
3237                         par_list_append(&ready_list, next_work_item);
3238         }
3239
3240         /*
3241          * main parent loop
3242          *
3243          * Keep going until there is no worker still running AND there is no work
3244          * left to be done.
3245          */
3246
3247         ahlog(AH, 1, "entering main parallel loop\n");
3248
3249         while ((next_work_item = get_next_work_item(AH, &ready_list,
3250                                                                                                 slots, n_slots)) != NULL ||
3251                    work_in_progress(slots, n_slots))
3252         {
3253                 if (next_work_item != NULL)
3254                 {
3255                         teReqs          reqs;
3256
3257                         /* If not to be dumped, don't waste time launching a worker */
3258                         reqs = _tocEntryRequired(next_work_item, AH->ropt, false);
3259                         if ((reqs & (REQ_SCHEMA | REQ_DATA)) == 0)
3260                         {
3261                                 ahlog(AH, 1, "skipping item %d %s %s\n",
3262                                           next_work_item->dumpId,
3263                                           next_work_item->desc, next_work_item->tag);
3264
3265                                 par_list_remove(next_work_item);
3266                                 reduce_dependencies(AH, next_work_item, &ready_list);
3267
3268                                 continue;
3269                         }
3270
3271                         if ((next_slot = get_next_slot(slots, n_slots)) != NO_SLOT)
3272                         {
3273                                 /* There is work still to do and a worker slot available */
3274                                 thandle         child;
3275                                 RestoreArgs *args;
3276
3277                                 ahlog(AH, 1, "launching item %d %s %s\n",
3278                                           next_work_item->dumpId,
3279                                           next_work_item->desc, next_work_item->tag);
3280
3281                                 par_list_remove(next_work_item);
3282
3283                                 /* this memory is dealloced in mark_work_done() */
3284                                 args = malloc(sizeof(RestoreArgs));
3285                                 args->AH = CloneArchive(AH);
3286                                 args->te = next_work_item;
3287
3288                                 /* run the step in a worker child */
3289                                 child = spawn_restore(args);
3290
3291                                 slots[next_slot].child_id = child;
3292                                 slots[next_slot].args = args;
3293
3294                                 continue;
3295                         }
3296                 }
3297
3298                 /*
3299                  * If we get here there must be work being done.  Either there is no
3300                  * work available to schedule (and work_in_progress returned true) or
3301                  * there are no slots available.  So we wait for a worker to finish,
3302                  * and process the result.
3303                  */
3304                 ret_child = reap_child(slots, n_slots, &work_status);
3305
3306                 if (WIFEXITED(work_status))
3307                 {
3308                         mark_work_done(AH, &ready_list,
3309                                                    ret_child, WEXITSTATUS(work_status),
3310                                                    slots, n_slots);
3311                 }
3312                 else
3313                 {
3314                         die_horribly(AH, modulename, "worker process crashed: status %d\n",
3315                                                  work_status);
3316                 }
3317         }
3318
3319         ahlog(AH, 1, "finished main parallel loop\n");
3320
3321         /*
3322          * Now reconnect the single parent connection.
3323          */
3324         ConnectDatabase((Archive *) AH, ropt->dbname,
3325                                         ropt->pghost, ropt->pgport, ropt->username,
3326                                         ropt->promptPassword);
3327
3328         _doSetFixedOutputState(AH);
3329
3330         /*
3331          * Make sure there is no non-ACL work left due to, say, circular
3332          * dependencies, or some other pathological condition. If so, do it in the
3333          * single parent connection.
3334          */
3335         for (te = pending_list.par_next; te != &pending_list; te = te->par_next)
3336         {
3337                 ahlog(AH, 1, "processing missed item %d %s %s\n",
3338                           te->dumpId, te->desc, te->tag);
3339                 (void) restore_toc_entry(AH, te, ropt, false);
3340         }
3341
3342         /* The ACLs will be handled back in RestoreArchive. */
3343 }
3344
3345 /*
3346  * create a worker child to perform a restore step in parallel
3347  */
3348 static thandle
3349 spawn_restore(RestoreArgs *args)
3350 {
3351         thandle         child;
3352
3353         /* Ensure stdio state is quiesced before forking */
3354         fflush(NULL);
3355
3356 #ifndef WIN32
3357         child = fork();
3358         if (child == 0)
3359         {
3360                 /* in child process */
3361                 parallel_restore(args);
3362                 die_horribly(args->AH, modulename,
3363                                          "parallel_restore should not return\n");
3364         }
3365         else if (child < 0)
3366         {
3367                 /* fork failed */
3368                 die_horribly(args->AH, modulename,
3369                                          "could not create worker process: %s\n",
3370                                          strerror(errno));
3371         }
3372 #else
3373         child = (HANDLE) _beginthreadex(NULL, 0, (void *) parallel_restore,
3374                                                                         args, 0, NULL);
3375         if (child == 0)
3376                 die_horribly(args->AH, modulename,
3377                                          "could not create worker thread: %s\n",
3378                                          strerror(errno));
3379 #endif
3380
3381         return child;
3382 }
3383
3384 /*
3385  *      collect status from a completed worker child
3386  */
3387 static thandle
3388 reap_child(ParallelSlot *slots, int n_slots, int *work_status)
3389 {
3390 #ifndef WIN32
3391         /* Unix is so much easier ... */
3392         return wait(work_status);
3393 #else
3394         static HANDLE *handles = NULL;
3395         int                     hindex,
3396                                 snum,
3397                                 tnum;
3398         thandle         ret_child;
3399         DWORD           res;
3400
3401         /* first time around only, make space for handles to listen on */
3402         if (handles == NULL)
3403                 handles = (HANDLE *) calloc(sizeof(HANDLE), n_slots);
3404
3405         /* set up list of handles to listen to */
3406         for (snum = 0, tnum = 0; snum < n_slots; snum++)
3407                 if (slots[snum].child_id != 0)
3408                         handles[tnum++] = slots[snum].child_id;
3409
3410         /* wait for one to finish */
3411         hindex = WaitForMultipleObjects(tnum, handles, false, INFINITE);
3412
3413         /* get handle of finished thread */
3414         ret_child = handles[hindex - WAIT_OBJECT_0];
3415
3416         /* get the result */
3417         GetExitCodeThread(ret_child, &res);
3418         *work_status = res;
3419
3420         /* dispose of handle to stop leaks */
3421         CloseHandle(ret_child);
3422
3423         return ret_child;
3424 #endif
3425 }
3426
3427 /*
3428  * are we doing anything now?
3429  */
3430 static bool
3431 work_in_progress(ParallelSlot *slots, int n_slots)
3432 {
3433         int                     i;
3434
3435         for (i = 0; i < n_slots; i++)
3436         {
3437                 if (slots[i].child_id != 0)
3438                         return true;
3439         }
3440         return false;
3441 }
3442
3443 /*
3444  * find the first free parallel slot (if any).
3445  */
3446 static int
3447 get_next_slot(ParallelSlot *slots, int n_slots)
3448 {
3449         int                     i;
3450
3451         for (i = 0; i < n_slots; i++)
3452         {
3453                 if (slots[i].child_id == 0)
3454                         return i;
3455         }
3456         return NO_SLOT;
3457 }
3458
3459
3460 /*
3461  * Check if te1 has an exclusive lock requirement for an item that te2 also
3462  * requires, whether or not te2's requirement is for an exclusive lock.
3463  */
3464 static bool
3465 has_lock_conflicts(TocEntry *te1, TocEntry *te2)
3466 {
3467         int                     j,
3468                                 k;
3469
3470         for (j = 0; j < te1->nLockDeps; j++)
3471         {
3472                 for (k = 0; k < te2->nDeps; k++)
3473                 {
3474                         if (te1->lockDeps[j] == te2->dependencies[k])
3475                                 return true;
3476                 }
3477         }
3478         return false;
3479 }
3480
3481
3482 /*
3483  * Initialize the header of a parallel-processing list.
3484  *
3485  * These are circular lists with a dummy TocEntry as header, just like the
3486  * main TOC list; but we use separate list links so that an entry can be in
3487  * the main TOC list as well as in a parallel-processing list.
3488  */
3489 static void
3490 par_list_header_init(TocEntry *l)
3491 {
3492         l->par_prev = l->par_next = l;
3493 }
3494
3495 /* Append te to the end of the parallel-processing list headed by l */
3496 static void
3497 par_list_append(TocEntry *l, TocEntry *te)
3498 {
3499         te->par_prev = l->par_prev;
3500         l->par_prev->par_next = te;
3501         l->par_prev = te;
3502         te->par_next = l;
3503 }
3504
3505 /* Remove te from whatever parallel-processing list it's in */
3506 static void
3507 par_list_remove(TocEntry *te)
3508 {
3509         te->par_prev->par_next = te->par_next;
3510         te->par_next->par_prev = te->par_prev;
3511         te->par_prev = NULL;
3512         te->par_next = NULL;
3513 }
3514
3515
3516 /*
3517  * Find the next work item (if any) that is capable of being run now.
3518  *
3519  * To qualify, the item must have no remaining dependencies
3520  * and no requirements for locks that are incompatible with
3521  * items currently running.  Items in the ready_list are known to have
3522  * no remaining dependencies, but we have to check for lock conflicts.
3523  *
3524  * Note that the returned item has *not* been removed from ready_list.
3525  * The caller must do that after successfully dispatching the item.
3526  *
3527  * pref_non_data is for an alternative selection algorithm that gives
3528  * preference to non-data items if there is already a data load running.
3529  * It is currently disabled.
3530  */
3531 static TocEntry *
3532 get_next_work_item(ArchiveHandle *AH, TocEntry *ready_list,
3533                                    ParallelSlot *slots, int n_slots)
3534 {
3535         bool            pref_non_data = false;  /* or get from AH->ropt */
3536         TocEntry   *data_te = NULL;
3537         TocEntry   *te;
3538         int                     i,
3539                                 k;
3540
3541         /*
3542          * Bogus heuristics for pref_non_data
3543          */
3544         if (pref_non_data)
3545         {
3546                 int                     count = 0;
3547
3548                 for (k = 0; k < n_slots; k++)
3549                         if (slots[k].args->te != NULL &&
3550                                 slots[k].args->te->section == SECTION_DATA)
3551                                 count++;
3552                 if (n_slots == 0 || count * 4 < n_slots)
3553                         pref_non_data = false;
3554         }
3555
3556         /*
3557          * Search the ready_list until we find a suitable item.
3558          */
3559         for (te = ready_list->par_next; te != ready_list; te = te->par_next)
3560         {
3561                 bool            conflicts = false;
3562
3563                 /*
3564                  * Check to see if the item would need exclusive lock on something
3565                  * that a currently running item also needs lock on, or vice versa. If
3566                  * so, we don't want to schedule them together.
3567                  */
3568                 for (i = 0; i < n_slots && !conflicts; i++)
3569                 {
3570                         TocEntry   *running_te;
3571
3572                         if (slots[i].args == NULL)
3573                                 continue;
3574                         running_te = slots[i].args->te;
3575
3576                         if (has_lock_conflicts(te, running_te) ||
3577                                 has_lock_conflicts(running_te, te))
3578                         {
3579                                 conflicts = true;
3580                                 break;
3581                         }
3582                 }
3583
3584                 if (conflicts)
3585                         continue;
3586
3587                 if (pref_non_data && te->section == SECTION_DATA)
3588                 {
3589                         if (data_te == NULL)
3590                                 data_te = te;
3591                         continue;
3592                 }
3593
3594                 /* passed all tests, so this item can run */
3595                 return te;
3596         }
3597
3598         if (data_te != NULL)
3599                 return data_te;
3600
3601         ahlog(AH, 2, "no item ready\n");
3602         return NULL;
3603 }
3604
3605
3606 /*
3607  * Restore a single TOC item in parallel with others
3608  *
3609  * this is the procedure run as a thread (Windows) or a
3610  * separate process (everything else).
3611  */
3612 static parallel_restore_result
3613 parallel_restore(RestoreArgs *args)
3614 {
3615         ArchiveHandle *AH = args->AH;
3616         TocEntry   *te = args->te;
3617         RestoreOptions *ropt = AH->ropt;
3618         int                     retval;
3619
3620         /*
3621          * Close and reopen the input file so we have a private file pointer that
3622          * doesn't stomp on anyone else's file pointer, if we're actually going to
3623          * need to read from the file. Otherwise, just close it except on Windows,
3624          * where it will possibly be needed by other threads.
3625          *
3626          * Note: on Windows, since we are using threads not processes, the reopen
3627          * call *doesn't* close the original file pointer but just open a new one.
3628          */
3629         if (te->section == SECTION_DATA)
3630                 (AH->ReopenPtr) (AH);
3631 #ifndef WIN32
3632         else
3633                 (AH->ClosePtr) (AH);
3634 #endif
3635
3636         /*
3637          * We need our own database connection, too
3638          */
3639         ConnectDatabase((Archive *) AH, ropt->dbname,
3640                                         ropt->pghost, ropt->pgport, ropt->username,
3641                                         ropt->promptPassword);
3642
3643         _doSetFixedOutputState(AH);
3644
3645         /* Restore the TOC item */
3646         retval = restore_toc_entry(AH, te, ropt, true);
3647
3648         /* And clean up */
3649         PQfinish(AH->connection);
3650         AH->connection = NULL;
3651
3652         /* If we reopened the file, we are done with it, so close it now */
3653         if (te->section == SECTION_DATA)
3654                 (AH->ClosePtr) (AH);
3655
3656         if (retval == 0 && AH->public.n_errors)
3657                 retval = WORKER_IGNORED_ERRORS;
3658
3659 #ifndef WIN32
3660         exit(retval);
3661 #else
3662         return retval;
3663 #endif
3664 }
3665
3666
3667 /*
3668  * Housekeeping to be done after a step has been parallel restored.
3669  *
3670  * Clear the appropriate slot, free all the extra memory we allocated,
3671  * update status, and reduce the dependency count of any dependent items.
3672  */
3673 static void
3674 mark_work_done(ArchiveHandle *AH, TocEntry *ready_list,
3675                            thandle worker, int status,
3676                            ParallelSlot *slots, int n_slots)
3677 {
3678         TocEntry   *te = NULL;
3679         int                     i;
3680
3681         for (i = 0; i < n_slots; i++)
3682         {
3683                 if (slots[i].child_id == worker)
3684                 {
3685                         slots[i].child_id = 0;
3686                         te = slots[i].args->te;
3687                         DeCloneArchive(slots[i].args->AH);
3688                         free(slots[i].args);
3689                         slots[i].args = NULL;
3690
3691                         break;
3692                 }
3693         }
3694
3695         if (te == NULL)
3696                 die_horribly(AH, modulename, "could not find slot of finished worker\n");
3697
3698         ahlog(AH, 1, "finished item %d %s %s\n",
3699                   te->dumpId, te->desc, te->tag);
3700
3701         if (status == WORKER_CREATE_DONE)
3702                 mark_create_done(AH, te);
3703         else if (status == WORKER_INHIBIT_DATA)
3704         {
3705                 inhibit_data_for_failed_table(AH, te);
3706                 AH->public.n_errors++;
3707         }
3708         else if (status == WORKER_IGNORED_ERRORS)
3709                 AH->public.n_errors++;
3710         else if (status != 0)
3711                 die_horribly(AH, modulename, "worker process failed: exit code %d\n",
3712                                          status);
3713
3714         reduce_dependencies(AH, te, ready_list);
3715 }
3716
3717
3718 /*
3719  * Process the dependency information into a form useful for parallel restore.
3720  *
3721  * We set up depCount fields that are the number of as-yet-unprocessed
3722  * dependencies for each TOC entry.
3723  *
3724  * We also identify locking dependencies so that we can avoid trying to
3725  * schedule conflicting items at the same time.
3726  */
3727 static void
3728 fix_dependencies(ArchiveHandle *AH)
3729 {
3730         TocEntry  **tocsByDumpId;
3731         TocEntry   *te;
3732         DumpId          maxDumpId;
3733         int                     i;
3734
3735         /*
3736          * For some of the steps here, it is convenient to have an array that
3737          * indexes the TOC entries by dump ID, rather than searching the TOC list
3738          * repeatedly.  Entries for dump IDs not present in the TOC will be NULL.
3739          *
3740          * NOTE: because maxDumpId is just the highest dump ID defined in the
3741          * archive, there might be dependencies for IDs > maxDumpId.  All uses of
3742          * this array must guard against out-of-range dependency numbers.
3743          *
3744          * Also, initialize the depCount fields, and make sure all the TOC items
3745          * are marked as not being in any parallel-processing list.
3746          */
3747         maxDumpId = AH->maxDumpId;
3748         tocsByDumpId = (TocEntry **) calloc(maxDumpId, sizeof(TocEntry *));
3749         for (te = AH->toc->next; te != AH->toc; te = te->next)
3750         {
3751                 tocsByDumpId[te->dumpId - 1] = te;
3752                 te->depCount = te->nDeps;
3753                 te->par_prev = NULL;
3754                 te->par_next = NULL;
3755         }
3756
3757         /*
3758          * POST_DATA items that are shown as depending on a table need to be
3759          * re-pointed to depend on that table's data, instead.  This ensures they
3760          * won't get scheduled until the data has been loaded.  We handle this by
3761          * first finding TABLE/TABLE DATA pairs and then scanning all the
3762          * dependencies.
3763          *
3764          * Note: currently, a TABLE DATA should always have exactly one
3765          * dependency, on its TABLE item.  So we don't bother to search, but look
3766          * just at the first dependency.  We do trouble to make sure that it's a
3767          * TABLE, if possible.  However, if the dependency isn't in the archive
3768          * then just assume it was a TABLE; this is to cover cases where the table
3769          * was suppressed but we have the data and some dependent post-data items.
3770          */
3771         for (te = AH->toc->next; te != AH->toc; te = te->next)
3772         {
3773                 if (strcmp(te->desc, "TABLE DATA") == 0 && te->nDeps > 0)
3774                 {
3775                         DumpId          tableId = te->dependencies[0];
3776
3777                         if (tableId > maxDumpId ||
3778                                 tocsByDumpId[tableId - 1] == NULL ||
3779                                 strcmp(tocsByDumpId[tableId - 1]->desc, "TABLE") == 0)
3780                         {
3781                                 repoint_table_dependencies(AH, tableId, te->dumpId);
3782                         }
3783                 }
3784         }
3785
3786         /*
3787          * Pre-8.4 versions of pg_dump neglected to set up a dependency from BLOB
3788          * COMMENTS to BLOBS.  Cope.  (We assume there's only one BLOBS and only
3789          * one BLOB COMMENTS in such files.)
3790          */
3791         if (AH->version < K_VERS_1_11)
3792         {
3793                 for (te = AH->toc->next; te != AH->toc; te = te->next)
3794                 {
3795                         if (strcmp(te->desc, "BLOB COMMENTS") == 0 && te->nDeps == 0)
3796                         {
3797                                 TocEntry   *te2;
3798
3799                                 for (te2 = AH->toc->next; te2 != AH->toc; te2 = te2->next)
3800                                 {
3801                                         if (strcmp(te2->desc, "BLOBS") == 0)
3802                                         {
3803                                                 te->dependencies = (DumpId *) malloc(sizeof(DumpId));
3804                                                 te->dependencies[0] = te2->dumpId;
3805                                                 te->nDeps++;
3806                                                 te->depCount++;
3807                                                 break;
3808                                         }
3809                                 }
3810                                 break;
3811                         }
3812                 }
3813         }
3814
3815         /*
3816          * It is possible that the dependencies list items that are not in the
3817          * archive at all.      Subtract such items from the depCounts.
3818          */
3819         for (te = AH->toc->next; te != AH->toc; te = te->next)
3820         {
3821                 for (i = 0; i < te->nDeps; i++)
3822                 {
3823                         DumpId          depid = te->dependencies[i];
3824
3825                         if (depid > maxDumpId || tocsByDumpId[depid - 1] == NULL)
3826                                 te->depCount--;
3827                 }
3828         }
3829
3830         /*
3831          * Lastly, work out the locking dependencies.
3832          */
3833         for (te = AH->toc->next; te != AH->toc; te = te->next)
3834         {
3835                 te->lockDeps = NULL;
3836                 te->nLockDeps = 0;
3837                 identify_locking_dependencies(te, tocsByDumpId, maxDumpId);
3838         }
3839
3840         free(tocsByDumpId);
3841 }
3842
3843 /*
3844  * Change dependencies on tableId to depend on tableDataId instead,
3845  * but only in POST_DATA items.
3846  */
3847 static void
3848 repoint_table_dependencies(ArchiveHandle *AH,
3849                                                    DumpId tableId, DumpId tableDataId)
3850 {
3851         TocEntry   *te;
3852         int                     i;
3853
3854         for (te = AH->toc->next; te != AH->toc; te = te->next)
3855         {
3856                 if (te->section != SECTION_POST_DATA)
3857                         continue;
3858                 for (i = 0; i < te->nDeps; i++)
3859                 {
3860                         if (te->dependencies[i] == tableId)
3861                         {
3862                                 te->dependencies[i] = tableDataId;
3863                                 ahlog(AH, 2, "transferring dependency %d -> %d to %d\n",
3864                                           te->dumpId, tableId, tableDataId);
3865                         }
3866                 }
3867         }
3868 }
3869
3870 /*
3871  * Identify which objects we'll need exclusive lock on in order to restore
3872  * the given TOC entry (*other* than the one identified by the TOC entry
3873  * itself).  Record their dump IDs in the entry's lockDeps[] array.
3874  * tocsByDumpId[] is a convenience array (of size maxDumpId) to avoid
3875  * searching the TOC for each dependency.
3876  */
3877 static void
3878 identify_locking_dependencies(TocEntry *te,
3879                                                           TocEntry **tocsByDumpId,
3880                                                           DumpId maxDumpId)
3881 {
3882         DumpId     *lockids;
3883         int                     nlockids;
3884         int                     i;
3885
3886         /* Quick exit if no dependencies at all */
3887         if (te->nDeps == 0)
3888                 return;
3889
3890         /* Exit if this entry doesn't need exclusive lock on other objects */
3891         if (!(strcmp(te->desc, "CONSTRAINT") == 0 ||
3892                   strcmp(te->desc, "CHECK CONSTRAINT") == 0 ||
3893                   strcmp(te->desc, "FK CONSTRAINT") == 0 ||
3894                   strcmp(te->desc, "RULE") == 0 ||
3895                   strcmp(te->desc, "TRIGGER") == 0))
3896                 return;
3897
3898         /*
3899          * We assume the item requires exclusive lock on each TABLE DATA item
3900          * listed among its dependencies.  (This was originally a dependency on
3901          * the TABLE, but fix_dependencies repointed it to the data item. Note
3902          * that all the entry types we are interested in here are POST_DATA, so
3903          * they will all have been changed this way.)
3904          */
3905         lockids = (DumpId *) malloc(te->nDeps * sizeof(DumpId));
3906         nlockids = 0;
3907         for (i = 0; i < te->nDeps; i++)
3908         {
3909                 DumpId          depid = te->dependencies[i];
3910
3911                 if (depid <= maxDumpId && tocsByDumpId[depid - 1] &&
3912                         strcmp(tocsByDumpId[depid - 1]->desc, "TABLE DATA") == 0)
3913                         lockids[nlockids++] = depid;
3914         }
3915
3916         if (nlockids == 0)
3917         {
3918                 free(lockids);
3919                 return;
3920         }
3921
3922         te->lockDeps = realloc(lockids, nlockids * sizeof(DumpId));
3923         te->nLockDeps = nlockids;
3924 }
3925
3926 /*
3927  * Remove the specified TOC entry from the depCounts of items that depend on
3928  * it, thereby possibly making them ready-to-run.  Any pending item that
3929  * becomes ready should be moved to the ready list.
3930  */
3931 static void
3932 reduce_dependencies(ArchiveHandle *AH, TocEntry *te, TocEntry *ready_list)
3933 {
3934         DumpId          target = te->dumpId;
3935         int                     i;
3936
3937         ahlog(AH, 2, "reducing dependencies for %d\n", target);
3938
3939         /*
3940          * We must examine all entries, not only the ones after the target item,
3941          * because if the user used a -L switch then the original dependency-
3942          * respecting order has been destroyed by SortTocFromFile.
3943          */
3944         for (te = AH->toc->next; te != AH->toc; te = te->next)
3945         {
3946                 for (i = 0; i < te->nDeps; i++)
3947                 {
3948                         if (te->dependencies[i] == target)
3949                         {
3950                                 te->depCount--;
3951                                 if (te->depCount == 0 && te->par_prev != NULL)
3952                                 {
3953                                         /* It must be in the pending list, so remove it ... */
3954                                         par_list_remove(te);
3955                                         /* ... and add to ready_list */
3956                                         par_list_append(ready_list, te);
3957                                 }
3958                         }
3959                 }
3960         }
3961 }
3962
3963 /*
3964  * Set the created flag on the DATA member corresponding to the given
3965  * TABLE member
3966  */
3967 static void
3968 mark_create_done(ArchiveHandle *AH, TocEntry *te)
3969 {
3970         TocEntry   *tes;
3971
3972         for (tes = AH->toc->next; tes != AH->toc; tes = tes->next)
3973         {
3974                 if (strcmp(tes->desc, "TABLE DATA") == 0 &&
3975                         strcmp(tes->tag, te->tag) == 0 &&
3976                         strcmp(tes->namespace ? tes->namespace : "",
3977                                    te->namespace ? te->namespace : "") == 0)
3978                 {
3979                         tes->created = true;
3980                         break;
3981                 }
3982         }
3983 }
3984
3985 /*
3986  * Mark the DATA member corresponding to the given TABLE member
3987  * as not wanted
3988  */
3989 static void
3990 inhibit_data_for_failed_table(ArchiveHandle *AH, TocEntry *te)
3991 {
3992         RestoreOptions *ropt = AH->ropt;
3993         TocEntry   *tes;
3994
3995         ahlog(AH, 1, "table \"%s\" could not be created, will not restore its data\n",
3996                   te->tag);
3997
3998         for (tes = AH->toc->next; tes != AH->toc; tes = tes->next)
3999         {
4000                 if (strcmp(tes->desc, "TABLE DATA") == 0 &&
4001                         strcmp(tes->tag, te->tag) == 0 &&
4002                         strcmp(tes->namespace ? tes->namespace : "",
4003                                    te->namespace ? te->namespace : "") == 0)
4004                 {
4005                         /* mark it unwanted; we assume idWanted array already exists */
4006                         ropt->idWanted[tes->dumpId - 1] = false;
4007                         break;
4008                 }
4009         }
4010 }
4011
4012
4013 /*
4014  * Clone and de-clone routines used in parallel restoration.
4015  *
4016  * Enough of the structure is cloned to ensure that there is no
4017  * conflict between different threads each with their own clone.
4018  *
4019  * These could be public, but no need at present.
4020  */
4021 static ArchiveHandle *
4022 CloneArchive(ArchiveHandle *AH)
4023 {
4024         ArchiveHandle *clone;
4025
4026         /* Make a "flat" copy */
4027         clone = (ArchiveHandle *) malloc(sizeof(ArchiveHandle));
4028         if (clone == NULL)
4029                 die_horribly(AH, modulename, "out of memory\n");
4030         memcpy(clone, AH, sizeof(ArchiveHandle));
4031
4032         /* Handle format-independent fields */
4033         clone->pgCopyBuf = createPQExpBuffer();
4034         clone->sqlBuf = createPQExpBuffer();
4035         clone->sqlparse.tagBuf = NULL;
4036
4037         /* The clone will have its own connection, so disregard connection state */
4038         clone->connection = NULL;
4039         clone->currUser = NULL;
4040         clone->currSchema = NULL;
4041         clone->currTablespace = NULL;
4042         clone->currWithOids = -1;
4043
4044         /* savedPassword must be local in case we change it while connecting */
4045         if (clone->savedPassword)
4046                 clone->savedPassword = strdup(clone->savedPassword);
4047
4048         /* clone has its own error count, too */
4049         clone->public.n_errors = 0;
4050
4051         /* Let the format-specific code have a chance too */
4052         (clone->ClonePtr) (clone);
4053
4054         return clone;
4055 }
4056
4057 /*
4058  * Release clone-local storage.
4059  *
4060  * Note: we assume any clone-local connection was already closed.
4061  */
4062 static void
4063 DeCloneArchive(ArchiveHandle *AH)
4064 {
4065         /* Clear format-specific state */
4066         (AH->DeClonePtr) (AH);
4067
4068         /* Clear state allocated by CloneArchive */
4069         destroyPQExpBuffer(AH->pgCopyBuf);
4070         destroyPQExpBuffer(AH->sqlBuf);
4071         if (AH->sqlparse.tagBuf)
4072                 destroyPQExpBuffer(AH->sqlparse.tagBuf);
4073
4074         /* Clear any connection-local state */
4075         if (AH->currUser)
4076                 free(AH->currUser);
4077         if (AH->currSchema)
4078                 free(AH->currSchema);
4079         if (AH->currTablespace)
4080                 free(AH->currTablespace);
4081         if (AH->savedPassword)
4082                 free(AH->savedPassword);
4083
4084         free(AH);
4085 }