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