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

   1 /*-------------------------------------------------------------------------
   2  *
   3  * pg_backup_archiver.c
   4  *
   5  *      Private implementation of the archiver routines.
   6  *
   7  *      See the headers to pg_restore for more details.
   8  *
   9  * Copyright (c) 2000, Philip Warner
  10  *      Rights are granted to use this software in any way so long
  11  *      as this notice is not removed.
  12  *
  13  *      The author is not responsible for loss or damages that may
  14  *      result from its use.
  15  *
  16  *
  17  * IDENTIFICATION
  18  *              $PostgreSQL: pgsql/src/bin/pg_dump/pg_backup_archiver.c,v 1.135 2006/08/01 18:21:44 momjian Exp $
  19  *
  20  *-------------------------------------------------------------------------
  21  */
  22
  23 #include "pg_backup_db.h"
  24 #include "dumputils.h"
  25
  26 #include <ctype.h>
  27
  28 #include <unistd.h>
  29
  30 #ifdef WIN32
  31 #include <io.h>
  32 #endif
  33
  34 #include "libpq/libpq-fs.h"
  35 #include "mb/pg_wchar.h"
  36
  37
  38 const char *progname;
  39
  40 static char *modulename = gettext_noop("archiver");
  41
  42
  43 static ArchiveHandle *_allocAH(const char *FileSpec, const ArchiveFormat fmt,
  44                  const int compression, ArchiveMode mode);
  45 static void _getObjectDescription(PQExpBuffer buf, TocEntry *te,
  46                                           ArchiveHandle *AH);
  47 static void _printTocEntry(ArchiveHandle *AH, TocEntry *te, RestoreOptions *ropt, bool isData, bool acl_pass);
  48
  49
  50 static void _doSetFixedOutputState(ArchiveHandle *AH);
  51 static void _doSetSessionAuth(ArchiveHandle *AH, const char *user);
  52 static void _doSetWithOids(ArchiveHandle *AH, const bool withOids);
  53 static void _reconnectToDB(ArchiveHandle *AH, const char *dbname);
  54 static void _becomeUser(ArchiveHandle *AH, const char *user);
  55 static void _becomeOwner(ArchiveHandle *AH, TocEntry *te);
  56 static void _selectOutputSchema(ArchiveHandle *AH, const char *schemaName);
  57 static void _selectTablespace(ArchiveHandle *AH, const char *tablespace);
  58 static void processEncodingEntry(ArchiveHandle *AH, TocEntry *te);
  59 static void processStdStringsEntry(ArchiveHandle *AH, TocEntry *te);
  60 static teReqs _tocEntryRequired(TocEntry *te, RestoreOptions *ropt, bool include_acls);
  61 static void _disableTriggersIfNecessary(ArchiveHandle *AH, TocEntry *te, RestoreOptions *ropt);
  62 static void _enableTriggersIfNecessary(ArchiveHandle *AH, TocEntry *te, RestoreOptions *ropt);
  63 static TocEntry *getTocEntryByDumpId(ArchiveHandle *AH, DumpId id);
  64 static void _moveAfter(ArchiveHandle *AH, TocEntry *pos, TocEntry *te);
  65 static int      _discoverArchiveFormat(ArchiveHandle *AH);
  66
  67 static void dump_lo_buf(ArchiveHandle *AH);
  68 static void _write_msg(const char *modulename, const char *fmt, va_list ap);
  69 static void _die_horribly(ArchiveHandle *AH, const char *modulename, const char *fmt, va_list ap);
  70
  71 static void dumpTimestamp(ArchiveHandle *AH, const char *msg, time_t tim);
  72 static OutputContext SetOutput(ArchiveHandle *AH, char *filename, int compression);
  73 static void ResetOutput(ArchiveHandle *AH, OutputContext savedContext);
  74
  75
  76 /*
  77  *      Wrapper functions.
  78  *
  79  *      The objective it to make writing new formats and dumpers as simple
  80  *      as possible, if necessary at the expense of extra function calls etc.
  81  *
  82  */
  83
  84
  85 /* Create a new archive */
  86 /* Public */
  87 Archive *
  88 CreateArchive(const char *FileSpec, const ArchiveFormat fmt,
  89                           const int compression)
  90
  91 {
  92         ArchiveHandle *AH = _allocAH(FileSpec, fmt, compression, archModeWrite);
  93
  94         return (Archive *) AH;
  95 }
  96
  97 /* Open an existing archive */
  98 /* Public */
  99 Archive *
 100 OpenArchive(const char *FileSpec, const ArchiveFormat fmt)
 101 {
 102         ArchiveHandle *AH = _allocAH(FileSpec, fmt, 0, archModeRead);
 103
 104         return (Archive *) AH;
 105 }
 106
 107 /* Public */
 108 void
 109 CloseArchive(Archive *AHX)
 110 {
 111         int                     res = 0;
 112         ArchiveHandle *AH = (ArchiveHandle *) AHX;
 113
 114         (*AH->ClosePtr) (AH);
 115
 116         /* Close the output */
 117         if (AH->gzOut)
 118                 res = GZCLOSE(AH->OF);
 119         else if (AH->OF != stdout)
 120                 res = fclose(AH->OF);
 121
 122         if (res != 0)
 123                 die_horribly(AH, modulename, "could not close output file: %s\n",
 124                                          strerror(errno));
 125 }
 126
 127 /* Public */
 128 void
 129 RestoreArchive(Archive *AHX, RestoreOptions *ropt)
 130 {
 131         ArchiveHandle *AH = (ArchiveHandle *) AHX;
 132         TocEntry   *te;
 133         teReqs          reqs;
 134         OutputContext sav;
 135         bool            defnDumped;
 136
 137         AH->ropt = ropt;
 138         AH->stage = STAGE_INITIALIZING;
 139
 140         /*
 141          * Check for nonsensical option combinations.
 142          *
 143          * NB: create+dropSchema is useless because if you're creating the DB,
 144          * there's no need to drop individual items in it.  Moreover, if we tried
 145          * to do that then we'd issue the drops in the database initially
 146          * connected to, not the one we will create, which is very bad...
 147          */
 148         if (ropt->create && ropt->dropSchema)
 149                 die_horribly(AH, modulename, "-C and -c are incompatible options\n");
 150
 151         /*
 152          * If we're using a DB connection, then connect it.
 153          */
 154         if (ropt->useDB)
 155         {
 156                 ahlog(AH, 1, "connecting to database for restore\n");
 157                 if (AH->version < K_VERS_1_3)
 158                         die_horribly(AH, modulename, "direct database connections are not supported in pre-1.3 archives\n");
 159
 160                 /* XXX Should get this from the archive */
 161                 AHX->minRemoteVersion = 070100;
 162                 AHX->maxRemoteVersion = 999999;
 163
 164                 ConnectDatabase(AHX, ropt->dbname,
 165                                                 ropt->pghost, ropt->pgport, ropt->username,
 166                                                 ropt->requirePassword, ropt->ignoreVersion);
 167
 168                 /*
 169                  * If we're talking to the DB directly, don't send comments since they
 170                  * obscure SQL when displaying errors
 171                  */
 172                 AH->noTocComments = 1;
 173         }
 174
 175         /*
 176          * Work out if we have an implied data-only restore. This can happen if
 177          * the dump was data only or if the user has used a toc list to exclude
 178          * all of the schema data. All we do is look for schema entries - if none
 179          * are found then we set the dataOnly flag.
 180          *
 181          * We could scan for wanted TABLE entries, but that is not the same as
 182          * dataOnly. At this stage, it seems unnecessary (6-Mar-2001).
 183          */
 184         if (!ropt->dataOnly)
 185         {
 186                 int                     impliedDataOnly = 1;
 187
 188                 for (te = AH->toc->next; te != AH->toc; te = te->next)
 189                 {
 190                         reqs = _tocEntryRequired(te, ropt, true);
 191                         if ((reqs & REQ_SCHEMA) != 0)
 192                         {                                       /* It's schema, and it's wanted */
 193                                 impliedDataOnly = 0;
 194                                 break;
 195                         }
 196                 }
 197                 if (impliedDataOnly)
 198                 {
 199                         ropt->dataOnly = impliedDataOnly;
 200                         ahlog(AH, 1, "implied data-only restore\n");
 201                 }
 202         }
 203
 204         /*
 205          * Setup the output file if necessary.
 206          */
 207         if (ropt->filename || ropt->compression)
 208                 sav = SetOutput(AH, ropt->filename, ropt->compression);
 209
 210         ahprintf(AH, "--\n-- PostgreSQL database dump\n--\n\n");
 211
 212         if (AH->public.verbose)
 213                 dumpTimestamp(AH, "Started on", AH->createDate);
 214
 215         if (ropt->single_txn)
 216         {
 217                 if (AH->connection)
 218                         StartTransaction(AH);
 219                 else
 220                         ahprintf(AH, "BEGIN;\n\n");
 221         }
 222
 223         /*
 224          * Establish important parameter values right away.
 225          */
 226         _doSetFixedOutputState(AH);
 227
 228         AH->stage = STAGE_PROCESSING;
 229
 230         /*
 231          * Drop the items at the start, in reverse order
 232          */
 233         if (ropt->dropSchema)
 234         {
 235                 for (te = AH->toc->prev; te != AH->toc; te = te->prev)
 236                 {
 237                         AH->currentTE = te;
 238
 239                         reqs = _tocEntryRequired(te, ropt, false /* needn't drop ACLs */ );
 240                         if (((reqs & REQ_SCHEMA) != 0) && te->dropStmt)
 241                         {
 242                                 /* We want the schema */
 243                                 ahlog(AH, 1, "dropping %s %s\n", te->desc, te->tag);
 244                                 /* Select owner and schema as necessary */
 245                                 _becomeOwner(AH, te);
 246                                 _selectOutputSchema(AH, te->namespace);
 247                                 /* Drop it */
 248                                 ahprintf(AH, "%s", te->dropStmt);
 249                         }
 250                 }
 251         }
 252
 253         /*
 254          * Now process each non-ACL TOC entry
 255          */
 256         for (te = AH->toc->next; te != AH->toc; te = te->next)
 257         {
 258                 AH->currentTE = te;
 259
 260                 /* Work out what, if anything, we want from this entry */
 261                 reqs = _tocEntryRequired(te, ropt, false);
 262
 263                 /* Dump any relevant dump warnings to stderr */
 264                 if (!ropt->suppressDumpWarnings && strcmp(te->desc, "WARNING") == 0)
 265                 {
 266                         if (!ropt->dataOnly && te->defn != NULL && strlen(te->defn) != 0)
 267                                 write_msg(modulename, "warning from original dump file: %s\n", te->defn);
 268                         else if (te->copyStmt != NULL && strlen(te->copyStmt) != 0)
 269                                 write_msg(modulename, "warning from original dump file: %s\n", te->copyStmt);
 270                 }
 271
 272                 defnDumped = false;
 273
 274                 if ((reqs & REQ_SCHEMA) != 0)   /* We want the schema */
 275                 {
 276                         ahlog(AH, 1, "creating %s %s\n", te->desc, te->tag);
 277
 278                         _printTocEntry(AH, te, ropt, false, false);
 279                         defnDumped = true;
 280
 281                         /* If we could not create a table, ignore the respective TABLE DATA if
 282                          * -X no-data-for-failed-tables is given */
 283                         if (ropt->noDataForFailedTables && AH->lastErrorTE == te && strcmp (te->desc, "TABLE") == 0) {
 284                                 TocEntry *tes, *last;
 285
 286                                 ahlog (AH, 1, "table %s could not be created, will not restore its data\n", te->tag);
 287
 288                                 for (last = te, tes = te->next; tes != AH->toc; last = tes, tes = tes->next) {
 289                                         if (strcmp (tes->desc, "TABLE DATA") == 0 && strcmp (tes->tag, te->tag) == 0 &&
 290                                             strcmp (tes->namespace ? tes->namespace : "", te->namespace ? te->namespace : "") == 0) {
 291                                             /* remove this node */
 292                                             last->next = tes->next;
 293                                             break;
 294                                         }
 295                                 }
 296                         }
 297
 298                         /* If we created a DB, connect to it... */
 299                         if (strcmp(te->desc, "DATABASE") == 0)
 300                         {
 301                                 ahlog(AH, 1, "connecting to new database \"%s\"\n", te->tag);
 302                                 _reconnectToDB(AH, te->tag);
 303                         }
 304                 }
 305
 306                 /*
 307                  * If we have a data component, then process it
 308                  */
 309                 if ((reqs & REQ_DATA) != 0)
 310                 {
 311                         /*
 312                          * hadDumper will be set if there is genuine data component for
 313                          * this node. Otherwise, we need to check the defn field for
 314                          * statements that need to be executed in data-only restores.
 315                          */
 316                         if (te->hadDumper)
 317                         {
 318                                 /*
 319                                  * If we can output the data, then restore it.
 320                                  */
 321                                 if (AH->PrintTocDataPtr !=NULL && (reqs & REQ_DATA) != 0)
 322                                 {
 323 #ifndef HAVE_LIBZ
 324                                         if (AH->compression != 0)
 325                                                 die_horribly(AH, modulename, "cannot restore from compressed archive (compression not supported in this installation)\n");
 326 #endif
 327
 328                                         _printTocEntry(AH, te, ropt, true, false);
 329
 330                                         if (strcmp(te->desc, "BLOBS") == 0 ||
 331                                                 strcmp(te->desc, "BLOB COMMENTS") == 0)
 332                                         {
 333                                                 ahlog(AH, 1, "restoring %s\n", te->desc);
 334
 335                                                 _selectOutputSchema(AH, "pg_catalog");
 336
 337                                                 (*AH->PrintTocDataPtr) (AH, te, ropt);
 338                                         }
 339                                         else
 340                                         {
 341                                                 _disableTriggersIfNecessary(AH, te, ropt);
 342
 343                                                 /* Select owner and schema as necessary */
 344                                                 _becomeOwner(AH, te);
 345                                                 _selectOutputSchema(AH, te->namespace);
 346
 347                                                 ahlog(AH, 1, "restoring data for table \"%s\"\n",
 348                                                           te->tag);
 349
 350                                                 /*
 351                                                  * If we have a copy statement, use it. As of V1.3,
 352                                                  * these are separate to allow easy import from
 353                                                  * withing a database connection. Pre 1.3 archives can
 354                                                  * not use DB connections and are sent to output only.
 355                                                  *
 356                                                  * For V1.3+, the table data MUST have a copy
 357                                                  * statement so that we can go into appropriate mode
 358                                                  * with libpq.
 359                                                  */
 360                                                 if (te->copyStmt && strlen(te->copyStmt) > 0)
 361                                                 {
 362                                                         ahprintf(AH, "%s", te->copyStmt);
 363                                                         AH->writingCopyData = true;
 364                                                 }
 365
 366                                                 (*AH->PrintTocDataPtr) (AH, te, ropt);
 367
 368                                                 AH->writingCopyData = false;
 369
 370                                                 _enableTriggersIfNecessary(AH, te, ropt);
 371                                         }
 372                                 }
 373                         }
 374                         else if (!defnDumped)
 375                         {
 376                                 /* If we haven't already dumped the defn part, do so now */
 377                                 ahlog(AH, 1, "executing %s %s\n", te->desc, te->tag);
 378                                 _printTocEntry(AH, te, ropt, false, false);
 379                         }
 380                 }
 381         }                                                       /* end loop over TOC entries */
 382
 383         /*
 384          * Scan TOC again to output ownership commands and ACLs
 385          */
 386         for (te = AH->toc->next; te != AH->toc; te = te->next)
 387         {
 388                 AH->currentTE = te;
 389
 390                 /* Work out what, if anything, we want from this entry */
 391                 reqs = _tocEntryRequired(te, ropt, true);
 392
 393                 if ((reqs & REQ_SCHEMA) != 0)   /* We want the schema */
 394                 {
 395                         ahlog(AH, 1, "setting owner and privileges for %s %s\n",
 396                                   te->desc, te->tag);
 397                         _printTocEntry(AH, te, ropt, false, true);
 398                 }
 399         }
 400
 401         if (ropt->single_txn)
 402         {
 403                 if (AH->connection)
 404                         CommitTransaction(AH);
 405                 else
 406                         ahprintf(AH, "COMMIT;\n\n");
 407         }
 408
 409         if (AH->public.verbose)
 410                 dumpTimestamp(AH, "Completed on", time(NULL));
 411
 412         ahprintf(AH, "--\n-- PostgreSQL database dump complete\n--\n\n");
 413
 414         /*
 415          * Clean up & we're done.
 416          */
 417         AH->stage = STAGE_FINALIZING;
 418
 419         if (ropt->filename || ropt->compression)
 420                 ResetOutput(AH, sav);
 421
 422         if (ropt->useDB)
 423         {
 424                 PQfinish(AH->connection);
 425                 AH->connection = NULL;
 426         }
 427 }
 428
 429 /*
 430  * Allocate a new RestoreOptions block.
 431  * This is mainly so we can initialize it, but also for future expansion,
 432  */
 433 RestoreOptions *
 434 NewRestoreOptions(void)
 435 {
 436         RestoreOptions *opts;
 437
 438         opts = (RestoreOptions *) calloc(1, sizeof(RestoreOptions));
 439
 440         opts->format = archUnknown;
 441         opts->suppressDumpWarnings = false;
 442         opts->exit_on_error = false;
 443
 444         return opts;
 445 }
 446
 447 static void
 448 _disableTriggersIfNecessary(ArchiveHandle *AH, TocEntry *te, RestoreOptions *ropt)
 449 {
 450         /* This hack is only needed in a data-only restore */
 451         if (!ropt->dataOnly || !ropt->disable_triggers)
 452                 return;
 453
 454         ahlog(AH, 1, "disabling triggers for %s\n", te->tag);
 455
 456         /*
 457          * Become superuser if possible, since they are the only ones who can
 458          * disable constraint triggers.  If -S was not given, assume the initial
 459          * user identity is a superuser.  (XXX would it be better to become the
 460          * table owner?)
 461          */
 462         _becomeUser(AH, ropt->superuser);
 463
 464         /*
 465          * Disable them.
 466          */
 467         _selectOutputSchema(AH, te->namespace);
 468
 469         ahprintf(AH, "ALTER TABLE %s DISABLE TRIGGER ALL;\n\n",
 470                          fmtId(te->tag));
 471 }
 472
 473 static void
 474 _enableTriggersIfNecessary(ArchiveHandle *AH, TocEntry *te, RestoreOptions *ropt)
 475 {
 476         /* This hack is only needed in a data-only restore */
 477         if (!ropt->dataOnly || !ropt->disable_triggers)
 478                 return;
 479
 480         ahlog(AH, 1, "enabling triggers for %s\n", te->tag);
 481
 482         /*
 483          * Become superuser if possible, since they are the only ones who can
 484          * disable constraint triggers.  If -S was not given, assume the initial
 485          * user identity is a superuser.  (XXX would it be better to become the
 486          * table owner?)
 487          */
 488         _becomeUser(AH, ropt->superuser);
 489
 490         /*
 491          * Enable them.
 492          */
 493         _selectOutputSchema(AH, te->namespace);
 494
 495         ahprintf(AH, "ALTER TABLE %s ENABLE TRIGGER ALL;\n\n",
 496                          fmtId(te->tag));
 497 }
 498
 499 /*
 500  * This is a routine that is part of the dumper interface, hence the 'Archive*' parameter.
 501  */
 502
 503 /* Public */
 504 size_t
 505 WriteData(Archive *AHX, const void *data, size_t dLen)
 506 {
 507         ArchiveHandle *AH = (ArchiveHandle *) AHX;
 508
 509         if (!AH->currToc)
 510                 die_horribly(AH, modulename, "internal error -- WriteData cannot be called outside the context of a DataDumper routine\n");
 511
 512         return (*AH->WriteDataPtr) (AH, data, dLen);
 513 }
 514
 515 /*
 516  * Create a new TOC entry. The TOC was designed as a TOC, but is now the
 517  * repository for all metadata. But the name has stuck.
 518  */
 519
 520 /* Public */
 521 void
 522 ArchiveEntry(Archive *AHX,
 523                          CatalogId catalogId, DumpId dumpId,
 524                          const char *tag,
 525                          const char *namespace,
 526                          const char *tablespace,
 527                          const char *owner, bool withOids,
 528                          const char *desc, const char *defn,
 529                          const char *dropStmt, const char *copyStmt,
 530                          const DumpId *deps, int nDeps,
 531                          DataDumperPtr dumpFn, void *dumpArg)
 532 {
 533         ArchiveHandle *AH = (ArchiveHandle *) AHX;
 534         TocEntry   *newToc;
 535
 536         newToc = (TocEntry *) calloc(1, sizeof(TocEntry));
 537         if (!newToc)
 538                 die_horribly(AH, modulename, "out of memory\n");
 539
 540         AH->tocCount++;
 541         if (dumpId > AH->maxDumpId)
 542                 AH->maxDumpId = dumpId;
 543
 544         newToc->prev = AH->toc->prev;
 545         newToc->next = AH->toc;
 546         AH->toc->prev->next = newToc;
 547         AH->toc->prev = newToc;
 548
 549         newToc->catalogId = catalogId;
 550         newToc->dumpId = dumpId;
 551
 552         newToc->tag = strdup(tag);
 553         newToc->namespace = namespace ? strdup(namespace) : NULL;
 554         newToc->tablespace = tablespace ? strdup(tablespace) : NULL;
 555         newToc->owner = strdup(owner);
 556         newToc->withOids = withOids;
 557         newToc->desc = strdup(desc);
 558         newToc->defn = strdup(defn);
 559         newToc->dropStmt = strdup(dropStmt);
 560         newToc->copyStmt = copyStmt ? strdup(copyStmt) : NULL;
 561
 562         if (nDeps > 0)
 563         {
 564                 newToc->dependencies = (DumpId *) malloc(nDeps * sizeof(DumpId));
 565                 memcpy(newToc->dependencies, deps, nDeps * sizeof(DumpId));
 566                 newToc->nDeps = nDeps;
 567         }
 568         else
 569         {
 570                 newToc->dependencies = NULL;
 571                 newToc->nDeps = 0;
 572         }
 573
 574         newToc->dataDumper = dumpFn;
 575         newToc->dataDumperArg = dumpArg;
 576         newToc->hadDumper = dumpFn ? true : false;
 577
 578         newToc->formatData = NULL;
 579
 580         if (AH->ArchiveEntryPtr !=NULL)
 581                 (*AH->ArchiveEntryPtr) (AH, newToc);
 582 }
 583
 584 /* Public */
 585 void
 586 PrintTOCSummary(Archive *AHX, RestoreOptions *ropt)
 587 {
 588         ArchiveHandle *AH = (ArchiveHandle *) AHX;
 589         TocEntry   *te = AH->toc->next;
 590         OutputContext sav;
 591         char       *fmtName;
 592
 593         if (ropt->filename)
 594                 sav = SetOutput(AH, ropt->filename, 0 /* no compression */ );
 595
 596         ahprintf(AH, ";\n; Archive created at %s", ctime(&AH->createDate));
 597         ahprintf(AH, ";     dbname: %s\n;     TOC Entries: %d\n;     Compression: %d\n",
 598                          AH->archdbname, AH->tocCount, AH->compression);
 599
 600         switch (AH->format)
 601         {
 602                 case archFiles:
 603                         fmtName = "FILES";
 604                         break;
 605                 case archCustom:
 606                         fmtName = "CUSTOM";
 607                         break;
 608                 case archTar:
 609                         fmtName = "TAR";
 610                         break;
 611                 default:
 612                         fmtName = "UNKNOWN";
 613         }
 614
 615         ahprintf(AH, ";     Dump Version: %d.%d-%d\n", AH->vmaj, AH->vmin, AH->vrev);
 616         ahprintf(AH, ";     Format: %s\n", fmtName);
 617         ahprintf(AH, ";     Integer: %d bytes\n", (int) AH->intSize);
 618         ahprintf(AH, ";     Offset: %d bytes\n", (int) AH->offSize);
 619         if (AH->archiveRemoteVersion)
 620                 ahprintf(AH, ";     Dumped from database version: %s\n",
 621                                  AH->archiveRemoteVersion);
 622         if (AH->archiveDumpVersion)
 623                 ahprintf(AH, ";     Dumped by pg_dump version: %s\n",
 624                                  AH->archiveDumpVersion);
 625
 626         ahprintf(AH, ";\n;\n; Selected TOC Entries:\n;\n");
 627
 628         while (te != AH->toc)
 629         {
 630                 if (_tocEntryRequired(te, ropt, true) != 0)
 631                         ahprintf(AH, "%d; %u %u %s %s %s %s\n", te->dumpId,
 632                                          te->catalogId.tableoid, te->catalogId.oid,
 633                                          te->desc, te->namespace ? te->namespace : "-",
 634                                          te->tag, te->owner);
 635                 te = te->next;
 636         }
 637
 638         if (ropt->filename)
 639                 ResetOutput(AH, sav);
 640 }
 641
 642 /***********
 643  * BLOB Archival
 644  ***********/
 645
 646 /* Called by a dumper to signal start of a BLOB */
 647 int
 648 StartBlob(Archive *AHX, Oid oid)
 649 {
 650         ArchiveHandle *AH = (ArchiveHandle *) AHX;
 651
 652         if (!AH->StartBlobPtr)
 653                 die_horribly(AH, modulename, "large-object output not supported in chosen format\n");
 654
 655         (*AH->StartBlobPtr) (AH, AH->currToc, oid);
 656
 657         return 1;
 658 }
 659
 660 /* Called by a dumper to signal end of a BLOB */
 661 int
 662 EndBlob(Archive *AHX, Oid oid)
 663 {
 664         ArchiveHandle *AH = (ArchiveHandle *) AHX;
 665
 666         if (AH->EndBlobPtr)
 667                 (*AH->EndBlobPtr) (AH, AH->currToc, oid);
 668
 669         return 1;
 670 }
 671
 672 /**********
 673  * BLOB Restoration
 674  **********/
 675
 676 /*
 677  * Called by a format handler before any blobs are restored
 678  */
 679 void
 680 StartRestoreBlobs(ArchiveHandle *AH)
 681 {
 682         if (!AH->ropt->single_txn)
 683         {
 684                 if (AH->connection)
 685                         StartTransaction(AH);
 686                 else
 687                         ahprintf(AH, "BEGIN;\n\n");
 688         }
 689
 690         AH->blobCount = 0;
 691 }
 692
 693 /*
 694  * Called by a format handler after all blobs are restored
 695  */
 696 void
 697 EndRestoreBlobs(ArchiveHandle *AH)
 698 {
 699         if (!AH->ropt->single_txn)
 700         {
 701                 if (AH->connection)
 702                         CommitTransaction(AH);
 703                 else
 704                         ahprintf(AH, "COMMIT;\n\n");
 705         }
 706
 707         ahlog(AH, 1, "restored %d large objects\n", AH->blobCount);
 708 }
 709
 710
 711 /*
 712  * Called by a format handler to initiate restoration of a blob
 713  */
 714 void
 715 StartRestoreBlob(ArchiveHandle *AH, Oid oid)
 716 {
 717         Oid                     loOid;
 718
 719         AH->blobCount++;
 720
 721         /* Initialize the LO Buffer */
 722         AH->lo_buf_used = 0;
 723
 724         ahlog(AH, 2, "restoring large object with OID %u\n", oid);
 725
 726         if (AH->connection)
 727         {
 728                 loOid = lo_create(AH->connection, oid);
 729                 if (loOid == 0 || loOid != oid)
 730                         die_horribly(AH, modulename, "could not create large object %u\n",
 731                                                  oid);
 732
 733                 AH->loFd = lo_open(AH->connection, oid, INV_WRITE);
 734                 if (AH->loFd == -1)
 735                         die_horribly(AH, modulename, "could not open large object\n");
 736         }
 737         else
 738         {
 739                 ahprintf(AH, "SELECT lo_open(lo_create(%u), %d);\n", oid, INV_WRITE);
 740         }
 741
 742         AH->writingBlob = 1;
 743 }
 744
 745 void
 746 EndRestoreBlob(ArchiveHandle *AH, Oid oid)
 747 {
 748         if (AH->lo_buf_used > 0)
 749         {
 750                 /* Write remaining bytes from the LO buffer */
 751                 dump_lo_buf(AH);
 752         }
 753
 754         AH->writingBlob = 0;
 755
 756         if (AH->connection)
 757         {
 758                 lo_close(AH->connection, AH->loFd);
 759                 AH->loFd = -1;
 760         }
 761         else
 762         {
 763                 ahprintf(AH, "SELECT lo_close(0);\n\n");
 764         }
 765 }
 766
 767 /***********
 768  * Sorting and Reordering
 769  ***********/
 770
 771 void
 772 SortTocFromFile(Archive *AHX, RestoreOptions *ropt)
 773 {
 774         ArchiveHandle *AH = (ArchiveHandle *) AHX;
 775         FILE       *fh;
 776         char            buf[1024];
 777         char       *cmnt;
 778         char       *endptr;
 779         DumpId          id;
 780         TocEntry   *te;
 781         TocEntry   *tePrev;
 782
 783         /* Allocate space for the 'wanted' array, and init it */
 784         ropt->idWanted = (bool *) malloc(sizeof(bool) * AH->maxDumpId);
 785         memset(ropt->idWanted, 0, sizeof(bool) * AH->maxDumpId);
 786         ropt->limitToList = true;
 787
 788         /* Set prev entry as head of list */
 789         tePrev = AH->toc;
 790
 791         /* Setup the file */
 792         fh = fopen(ropt->tocFile, PG_BINARY_R);
 793         if (!fh)
 794                 die_horribly(AH, modulename, "could not open TOC file: %s\n",
 795                                          strerror(errno));
 796
 797         while (fgets(buf, sizeof(buf), fh) != NULL)
 798         {
 799                 /* Truncate line at comment, if any */
 800                 cmnt = strchr(buf, ';');
 801                 if (cmnt != NULL)
 802                         cmnt[0] = '\0';
 803
 804                 /* Ignore if all blank */
 805                 if (strspn(buf, " \t\r") == strlen(buf))
 806                         continue;
 807
 808                 /* Get an ID, check it's valid and not already seen */
 809                 id = strtol(buf, &endptr, 10);
 810                 if (endptr == buf || id <= 0 || id > AH->maxDumpId ||
 811                         ropt->idWanted[id - 1])
 812                 {
 813                         write_msg(modulename, "WARNING: line ignored: %s\n", buf);
 814                         continue;
 815                 }
 816
 817                 /* Find TOC entry */
 818                 te = getTocEntryByDumpId(AH, id);
 819                 if (!te)
 820                         die_horribly(AH, modulename, "could not find entry for ID %d\n",
 821                                                  id);
 822
 823                 ropt->idWanted[id - 1] = true;
 824
 825                 _moveAfter(AH, tePrev, te);
 826                 tePrev = te;
 827         }
 828
 829         if (fclose(fh) != 0)
 830                 die_horribly(AH, modulename, "could not close TOC file: %s\n",
 831                                          strerror(errno));
 832 }
 833
 834 /**********************
 835  * 'Convenience functions that look like standard IO functions
 836  * for writing data when in dump mode.
 837  **********************/
 838
 839 /* Public */
 840 int
 841 archputs(const char *s, Archive *AH)
 842 {
 843         return WriteData(AH, s, strlen(s));
 844 }
 845
 846 /* Public */
 847 int
 848 archprintf(Archive *AH, const char *fmt,...)
 849 {
 850         char       *p = NULL;
 851         va_list         ap;
 852         int                     bSize = strlen(fmt) + 256;
 853         int                     cnt = -1;
 854
 855         /*
 856          * This is paranoid: deal with the possibility that vsnprintf is willing
 857          * to ignore trailing null or returns > 0 even if string does not fit. It
 858          * may be the case that it returns cnt = bufsize
 859          */
 860         while (cnt < 0 || cnt >= (bSize - 1))
 861         {
 862                 if (p != NULL)
 863                         free(p);
 864                 bSize *= 2;
 865                 p = (char *) malloc(bSize);
 866                 if (p == NULL)
 867                         exit_horribly(AH, modulename, "out of memory\n");
 868                 va_start(ap, fmt);
 869                 cnt = vsnprintf(p, bSize, fmt, ap);
 870                 va_end(ap);
 871         }
 872         WriteData(AH, p, cnt);
 873         free(p);
 874         return cnt;
 875 }
 876
 877
 878 /*******************************
 879  * Stuff below here should be 'private' to the archiver routines
 880  *******************************/
 881
 882 static OutputContext
 883 SetOutput(ArchiveHandle *AH, char *filename, int compression)
 884 {
 885         OutputContext sav;
 886         int                     fn;
 887
 888         /* Replace the AH output file handle */
 889         sav.OF = AH->OF;
 890         sav.gzOut = AH->gzOut;
 891
 892         if (filename)
 893                 fn = -1;
 894         else if (AH->FH)
 895                 fn = fileno(AH->FH);
 896         else if (AH->fSpec)
 897         {
 898                 fn = -1;
 899                 filename = AH->fSpec;
 900         }
 901         else
 902                 fn = fileno(stdout);
 903
 904         /* If compression explicitly requested, use gzopen */
 905 #ifdef HAVE_LIBZ
 906         if (compression != 0)
 907         {
 908                 char            fmode[10];
 909
 910                 /* Don't use PG_BINARY_x since this is zlib */
 911                 sprintf(fmode, "wb%d", compression);
 912                 if (fn >= 0)
 913                         AH->OF = gzdopen(dup(fn), fmode);
 914                 else
 915                         AH->OF = gzopen(filename, fmode);
 916                 AH->gzOut = 1;
 917         }
 918         else
 919 #endif
 920         {                                                       /* Use fopen */
 921                 if (fn >= 0)
 922                         AH->OF = fdopen(dup(fn), PG_BINARY_W);
 923                 else
 924                         AH->OF = fopen(filename, PG_BINARY_W);
 925                 AH->gzOut = 0;
 926         }
 927
 928         if (!AH->OF)
 929                 die_horribly(AH, modulename, "could not open output file: %s\n", strerror(errno));
 930
 931         return sav;
 932 }
 933
 934 static void
 935 ResetOutput(ArchiveHandle *AH, OutputContext sav)
 936 {
 937         int                     res;
 938
 939         if (AH->gzOut)
 940                 res = GZCLOSE(AH->OF);
 941         else
 942                 res = fclose(AH->OF);
 943
 944         if (res != 0)
 945                 die_horribly(AH, modulename, "could not close output file: %s\n",
 946                                          strerror(errno));
 947
 948         AH->gzOut = sav.gzOut;
 949         AH->OF = sav.OF;
 950 }
 951
 952
 953
 954 /*
 955  *      Print formatted text to the output file (usually stdout).
 956  */
 957 int
 958 ahprintf(ArchiveHandle *AH, const char *fmt,...)
 959 {
 960         char       *p = NULL;
 961         va_list         ap;
 962         int                     bSize = strlen(fmt) + 256;              /* Should be enough */
 963         int                     cnt = -1;
 964
 965         /*
 966          * This is paranoid: deal with the possibility that vsnprintf is willing
 967          * to ignore trailing null
 968          */
 969
 970         /*
 971          * or returns > 0 even if string does not fit. It may be the case that it
 972          * returns cnt = bufsize
 973          */
 974         while (cnt < 0 || cnt >= (bSize - 1))
 975         {
 976                 if (p != NULL)
 977                         free(p);
 978                 bSize *= 2;
 979                 p = (char *) malloc(bSize);
 980                 if (p == NULL)
 981                         die_horribly(AH, modulename, "out of memory\n");
 982                 va_start(ap, fmt);
 983                 cnt = vsnprintf(p, bSize, fmt, ap);
 984                 va_end(ap);
 985         }
 986         ahwrite(p, 1, cnt, AH);
 987         free(p);
 988         return cnt;
 989 }
 990
 991 void
 992 ahlog(ArchiveHandle *AH, int level, const char *fmt,...)
 993 {
 994         va_list         ap;
 995
 996         if (AH->debugLevel < level && (!AH->public.verbose || level > 1))
 997                 return;
 998
 999         va_start(ap, fmt);
1000         _write_msg(NULL, fmt, ap);
1001         va_end(ap);
1002 }
1003
1004 /*
1005  * Single place for logic which says 'We are restoring to a direct DB connection'.
1006  */
1007 static int
1008 RestoringToDB(ArchiveHandle *AH)
1009 {
1010         return (AH->ropt && AH->ropt->useDB && AH->connection);
1011 }
1012
1013 /*
1014  * Dump the current contents of the LO data buffer while writing a BLOB
1015  */
1016 static void
1017 dump_lo_buf(ArchiveHandle *AH)
1018 {
1019         if (AH->connection)
1020         {
1021                 size_t          res;
1022
1023                 res = lo_write(AH->connection, AH->loFd, AH->lo_buf, AH->lo_buf_used);
1024                 ahlog(AH, 5, "wrote %lu bytes of large object data (result = %lu)\n",
1025                           (unsigned long) AH->lo_buf_used, (unsigned long) res);
1026                 if (res != AH->lo_buf_used)
1027                         die_horribly(AH, modulename,
1028                         "could not write to large object (result: %lu, expected: %lu)\n",
1029                                            (unsigned long) res, (unsigned long) AH->lo_buf_used);
1030         }
1031         else
1032         {
1033                 unsigned char *str;
1034                 size_t          len;
1035
1036                 str = PQescapeBytea((const unsigned char *) AH->lo_buf,
1037                                                         AH->lo_buf_used, &len);
1038                 if (!str)
1039                         die_horribly(AH, modulename, "out of memory\n");
1040
1041                 /* Hack: turn off writingBlob so ahwrite doesn't recurse to here */
1042                 AH->writingBlob = 0;
1043                 ahprintf(AH, "SELECT lowrite(0, '%s');\n", str);
1044                 AH->writingBlob = 1;
1045
1046                 free(str);
1047         }
1048         AH->lo_buf_used = 0;
1049 }
1050
1051
1052 /*
1053  *      Write buffer to the output file (usually stdout). This is user for
1054  *      outputting 'restore' scripts etc. It is even possible for an archive
1055  *      format to create a custom output routine to 'fake' a restore if it
1056  *      wants to generate a script (see TAR output).
1057  */
1058 int
1059 ahwrite(const void *ptr, size_t size, size_t nmemb, ArchiveHandle *AH)
1060 {
1061         size_t          res;
1062
1063         if (AH->writingBlob)
1064         {
1065                 size_t          remaining = size * nmemb;
1066
1067                 while (AH->lo_buf_used + remaining > AH->lo_buf_size)
1068                 {
1069                         size_t          avail = AH->lo_buf_size - AH->lo_buf_used;
1070
1071                         memcpy((char *) AH->lo_buf + AH->lo_buf_used, ptr, avail);
1072                         ptr = (const void *) ((const char *) ptr + avail);
1073                         remaining -= avail;
1074                         AH->lo_buf_used += avail;
1075                         dump_lo_buf(AH);
1076                 }
1077
1078                 memcpy((char *) AH->lo_buf + AH->lo_buf_used, ptr, remaining);
1079                 AH->lo_buf_used += remaining;
1080
1081                 return size * nmemb;
1082         }
1083         else if (AH->gzOut)
1084         {
1085                 res = GZWRITE((void *) ptr, size, nmemb, AH->OF);
1086                 if (res != (nmemb * size))
1087                         die_horribly(AH, modulename, "could not write to output file: %s\n", strerror(errno));
1088                 return res;
1089         }
1090         else if (AH->CustomOutPtr)
1091         {
1092                 res = AH->CustomOutPtr (AH, ptr, size * nmemb);
1093
1094                 if (res != (nmemb * size))
1095                         die_horribly(AH, modulename, "could not write to custom output routine\n");
1096                 return res;
1097         }
1098         else
1099         {
1100                 /*
1101                  * If we're doing a restore, and it's direct to DB, and we're
1102                  * connected then send it to the DB.
1103                  */
1104                 if (RestoringToDB(AH))
1105                         return ExecuteSqlCommandBuf(AH, (void *) ptr, size * nmemb);            /* Always 1, currently */
1106                 else
1107                 {
1108                         res = fwrite((void *) ptr, size, nmemb, AH->OF);
1109                         if (res != nmemb)
1110                                 die_horribly(AH, modulename, "could not write to output file: %s\n",
1111                                                          strerror(errno));
1112                         return res;
1113                 }
1114         }
1115 }
1116
1117 /* Common exit code */
1118 static void
1119 _write_msg(const char *modulename, const char *fmt, va_list ap)
1120 {
1121         if (modulename)
1122                 fprintf(stderr, "%s: [%s] ", progname, _(modulename));
1123         else
1124                 fprintf(stderr, "%s: ", progname);
1125         vfprintf(stderr, _(fmt), ap);
1126 }
1127
1128 void
1129 write_msg(const char *modulename, const char *fmt,...)
1130 {
1131         va_list         ap;
1132
1133         va_start(ap, fmt);
1134         _write_msg(modulename, fmt, ap);
1135         va_end(ap);
1136 }
1137
1138
1139 static void
1140 _die_horribly(ArchiveHandle *AH, const char *modulename, const char *fmt, va_list ap)
1141 {
1142         _write_msg(modulename, fmt, ap);
1143
1144         if (AH)
1145         {
1146                 if (AH->public.verbose)
1147                         write_msg(NULL, "*** aborted because of error\n");
1148                 if (AH->connection)
1149                         PQfinish(AH->connection);
1150         }
1151
1152         exit(1);
1153 }
1154
1155 /* External use */
1156 void
1157 exit_horribly(Archive *AH, const char *modulename, const char *fmt,...)
1158 {
1159         va_list         ap;
1160
1161         va_start(ap, fmt);
1162         _die_horribly((ArchiveHandle *) AH, modulename, fmt, ap);
1163         va_end(ap);
1164 }
1165
1166 /* Archiver use (just different arg declaration) */
1167 void
1168 die_horribly(ArchiveHandle *AH, const char *modulename, const char *fmt,...)
1169 {
1170         va_list         ap;
1171
1172         va_start(ap, fmt);
1173         _die_horribly(AH, modulename, fmt, ap);
1174         va_end(ap);
1175 }
1176
1177 /* on some error, we may decide to go on... */
1178 void
1179 warn_or_die_horribly(ArchiveHandle *AH,
1180                                          const char *modulename, const char *fmt,...)
1181 {
1182         va_list         ap;
1183
1184         switch (AH->stage)
1185         {
1186
1187                 case STAGE_NONE:
1188                         /* Do nothing special */
1189                         break;
1190
1191                 case STAGE_INITIALIZING:
1192                         if (AH->stage != AH->lastErrorStage)
1193                                 write_msg(modulename, "Error while INITIALIZING:\n");
1194                         break;
1195
1196                 case STAGE_PROCESSING:
1197                         if (AH->stage != AH->lastErrorStage)
1198                                 write_msg(modulename, "Error while PROCESSING TOC:\n");
1199                         break;
1200
1201                 case STAGE_FINALIZING:
1202                         if (AH->stage != AH->lastErrorStage)
1203                                 write_msg(modulename, "Error while FINALIZING:\n");
1204                         break;
1205         }
1206         if (AH->currentTE != NULL && AH->currentTE != AH->lastErrorTE)
1207         {
1208                 write_msg(modulename, "Error from TOC entry %d; %u %u %s %s %s\n",
1209                                   AH->currentTE->dumpId,
1210                          AH->currentTE->catalogId.tableoid, AH->currentTE->catalogId.oid,
1211                           AH->currentTE->desc, AH->currentTE->tag, AH->currentTE->owner);
1212         }
1213         AH->lastErrorStage = AH->stage;
1214         AH->lastErrorTE = AH->currentTE;
1215
1216         va_start(ap, fmt);
1217         if (AH->public.exit_on_error)
1218                 _die_horribly(AH, modulename, fmt, ap);
1219         else
1220         {
1221                 _write_msg(modulename, fmt, ap);
1222                 AH->public.n_errors++;
1223         }
1224         va_end(ap);
1225 }
1226
1227 static void
1228 _moveAfter(ArchiveHandle *AH, TocEntry *pos, TocEntry *te)
1229 {
1230         te->prev->next = te->next;
1231         te->next->prev = te->prev;
1232
1233         te->prev = pos;
1234         te->next = pos->next;
1235
1236         pos->next->prev = te;
1237         pos->next = te;
1238 }
1239
1240 #ifdef NOT_USED
1241
1242 static void
1243 _moveBefore(ArchiveHandle *AH, TocEntry *pos, TocEntry *te)
1244 {
1245         te->prev->next = te->next;
1246         te->next->prev = te->prev;
1247
1248         te->prev = pos->prev;
1249         te->next = pos;
1250         pos->prev->next = te;
1251         pos->prev = te;
1252 }
1253 #endif
1254
1255 static TocEntry *
1256 getTocEntryByDumpId(ArchiveHandle *AH, DumpId id)
1257 {
1258         TocEntry   *te;
1259
1260         te = AH->toc->next;
1261         while (te != AH->toc)
1262         {
1263                 if (te->dumpId == id)
1264                         return te;
1265                 te = te->next;
1266         }
1267         return NULL;
1268 }
1269
1270 teReqs
1271 TocIDRequired(ArchiveHandle *AH, DumpId id, RestoreOptions *ropt)
1272 {
1273         TocEntry   *te = getTocEntryByDumpId(AH, id);
1274
1275         if (!te)
1276                 return 0;
1277
1278         return _tocEntryRequired(te, ropt, true);
1279 }
1280
1281 size_t
1282 WriteOffset(ArchiveHandle *AH, off_t o, int wasSet)
1283 {
1284         int                     off;
1285
1286         /* Save the flag */
1287         (*AH->WriteBytePtr) (AH, wasSet);
1288
1289         /* Write out off_t smallest byte first, prevents endian mismatch */
1290         for (off = 0; off < sizeof(off_t); off++)
1291         {
1292                 (*AH->WriteBytePtr) (AH, o & 0xFF);
1293                 o >>= 8;
1294         }
1295         return sizeof(off_t) + 1;
1296 }
1297
1298 int
1299 ReadOffset(ArchiveHandle *AH, off_t *o)
1300 {
1301         int                     i;
1302         int                     off;
1303         int                     offsetFlg;
1304
1305         /* Initialize to zero */
1306         *o = 0;
1307
1308         /* Check for old version */
1309         if (AH->version < K_VERS_1_7)
1310         {
1311                 /* Prior versions wrote offsets using WriteInt */
1312                 i = ReadInt(AH);
1313                 /* -1 means not set */
1314                 if (i < 0)
1315                         return K_OFFSET_POS_NOT_SET;
1316                 else if (i == 0)
1317                         return K_OFFSET_NO_DATA;
1318
1319                 /* Cast to off_t because it was written as an int. */
1320                 *o = (off_t) i;
1321                 return K_OFFSET_POS_SET;
1322         }
1323
1324         /*
1325          * Read the flag indicating the state of the data pointer. Check if valid
1326          * and die if not.
1327          *
1328          * This used to be handled by a negative or zero pointer, now we use an
1329          * extra byte specifically for the state.
1330          */
1331         offsetFlg = (*AH->ReadBytePtr) (AH) & 0xFF;
1332
1333         switch (offsetFlg)
1334         {
1335                 case K_OFFSET_POS_NOT_SET:
1336                 case K_OFFSET_NO_DATA:
1337                 case K_OFFSET_POS_SET:
1338
1339                         break;
1340
1341                 default:
1342                         die_horribly(AH, modulename, "unexpected data offset flag %d\n", offsetFlg);
1343         }
1344
1345         /*
1346          * Read the bytes
1347          */
1348         for (off = 0; off < AH->offSize; off++)
1349         {
1350                 if (off < sizeof(off_t))
1351                         *o |= ((off_t) ((*AH->ReadBytePtr) (AH))) << (off * 8);
1352                 else
1353                 {
1354                         if ((*AH->ReadBytePtr) (AH) != 0)
1355                                 die_horribly(AH, modulename, "file offset in dump file is too large\n");
1356                 }
1357         }
1358
1359         return offsetFlg;
1360 }
1361
1362 size_t
1363 WriteInt(ArchiveHandle *AH, int i)
1364 {
1365         int                     b;
1366
1367         /*
1368          * This is a bit yucky, but I don't want to make the binary format very
1369          * dependent on representation, and not knowing much about it, I write out
1370          * a sign byte. If you change this, don't forget to change the file
1371          * version #, and modify readInt to read the new format AS WELL AS the old
1372          * formats.
1373          */
1374
1375         /* SIGN byte */
1376         if (i < 0)
1377         {
1378                 (*AH->WriteBytePtr) (AH, 1);
1379                 i = -i;
1380         }
1381         else
1382                 (*AH->WriteBytePtr) (AH, 0);
1383
1384         for (b = 0; b < AH->intSize; b++)
1385         {
1386                 (*AH->WriteBytePtr) (AH, i & 0xFF);
1387                 i >>= 8;
1388         }
1389
1390         return AH->intSize + 1;
1391 }
1392
1393 int
1394 ReadInt(ArchiveHandle *AH)
1395 {
1396         int                     res = 0;
1397         int                     bv,
1398                                 b;
1399         int                     sign = 0;               /* Default positive */
1400         int                     bitShift = 0;
1401
1402         if (AH->version > K_VERS_1_0)
1403                 /* Read a sign byte */
1404                 sign = (*AH->ReadBytePtr) (AH);
1405
1406         for (b = 0; b < AH->intSize; b++)
1407         {
1408                 bv = (*AH->ReadBytePtr) (AH) & 0xFF;
1409                 if (bv != 0)
1410                         res = res + (bv << bitShift);
1411                 bitShift += 8;
1412         }
1413
1414         if (sign)
1415                 res = -res;
1416
1417         return res;
1418 }
1419
1420 size_t
1421 WriteStr(ArchiveHandle *AH, const char *c)
1422 {
1423         size_t          res;
1424
1425         if (c)
1426         {
1427                 res = WriteInt(AH, strlen(c));
1428                 res += (*AH->WriteBufPtr) (AH, c, strlen(c));
1429         }
1430         else
1431                 res = WriteInt(AH, -1);
1432
1433         return res;
1434 }
1435
1436 char *
1437 ReadStr(ArchiveHandle *AH)
1438 {
1439         char       *buf;
1440         int                     l;
1441
1442         l = ReadInt(AH);
1443         if (l == -1)
1444                 buf = NULL;
1445         else
1446         {
1447                 buf = (char *) malloc(l + 1);
1448                 if (!buf)
1449                         die_horribly(AH, modulename, "out of memory\n");
1450
1451                 (*AH->ReadBufPtr) (AH, (void *) buf, l);
1452                 buf[l] = '\0';
1453         }
1454
1455         return buf;
1456 }
1457
1458 static int
1459 _discoverArchiveFormat(ArchiveHandle *AH)
1460 {
1461         FILE       *fh;
1462         char            sig[6];                 /* More than enough */
1463         size_t          cnt;
1464         int                     wantClose = 0;
1465
1466 #if 0
1467         write_msg(modulename, "attempting to ascertain archive format\n");
1468 #endif
1469
1470         if (AH->lookahead)
1471                 free(AH->lookahead);
1472
1473         AH->lookaheadSize = 512;
1474         AH->lookahead = calloc(1, 512);
1475         AH->lookaheadLen = 0;
1476         AH->lookaheadPos = 0;
1477
1478         if (AH->fSpec)
1479         {
1480                 wantClose = 1;
1481                 fh = fopen(AH->fSpec, PG_BINARY_R);
1482         }
1483         else
1484                 fh = stdin;
1485
1486         if (!fh)
1487                 die_horribly(AH, modulename, "could not open input file: %s\n", strerror(errno));
1488
1489         cnt = fread(sig, 1, 5, fh);
1490
1491         if (cnt != 5)
1492         {
1493                 if (ferror(fh))
1494                         die_horribly(AH, modulename, "could not read input file: %s\n", strerror(errno));
1495                 else
1496                         die_horribly(AH, modulename, "input file is too short (read %lu, expected 5)\n",
1497                                                  (unsigned long) cnt);
1498         }
1499
1500         /* Save it, just in case we need it later */
1501         strncpy(&AH->lookahead[0], sig, 5);
1502         AH->lookaheadLen = 5;
1503
1504         if (strncmp(sig, "PGDMP", 5) == 0)
1505         {
1506                 AH->vmaj = fgetc(fh);
1507                 AH->vmin = fgetc(fh);
1508
1509                 /* Save these too... */
1510                 AH->lookahead[AH->lookaheadLen++] = AH->vmaj;
1511                 AH->lookahead[AH->lookaheadLen++] = AH->vmin;
1512
1513                 /* Check header version; varies from V1.0 */
1514                 if (AH->vmaj > 1 || ((AH->vmaj == 1) && (AH->vmin > 0)))                /* Version > 1.0 */
1515                 {
1516                         AH->vrev = fgetc(fh);
1517                         AH->lookahead[AH->lookaheadLen++] = AH->vrev;
1518                 }
1519                 else
1520                         AH->vrev = 0;
1521
1522                 /* Make a convenient integer <maj><min><rev>00 */
1523                 AH->version = ((AH->vmaj * 256 + AH->vmin) * 256 + AH->vrev) * 256 + 0;
1524
1525                 AH->intSize = fgetc(fh);
1526                 AH->lookahead[AH->lookaheadLen++] = AH->intSize;
1527
1528                 if (AH->version >= K_VERS_1_7)
1529                 {
1530                         AH->offSize = fgetc(fh);
1531                         AH->lookahead[AH->lookaheadLen++] = AH->offSize;
1532                 }
1533                 else
1534                         AH->offSize = AH->intSize;
1535
1536                 AH->format = fgetc(fh);
1537                 AH->lookahead[AH->lookaheadLen++] = AH->format;
1538         }
1539         else
1540         {
1541                 /*
1542                  * *Maybe* we have a tar archive format file... So, read first 512
1543                  * byte header...
1544                  */
1545                 cnt = fread(&AH->lookahead[AH->lookaheadLen], 1, 512 - AH->lookaheadLen, fh);
1546                 AH->lookaheadLen += cnt;
1547
1548                 if (AH->lookaheadLen != 512)
1549                         die_horribly(AH, modulename, "input file does not appear to be a valid archive (too short?)\n");
1550
1551                 if (!isValidTarHeader(AH->lookahead))
1552                         die_horribly(AH, modulename, "input file does not appear to be a valid archive\n");
1553
1554                 AH->format = archTar;
1555         }
1556
1557         /* If we can't seek, then mark the header as read */
1558         if (fseeko(fh, 0, SEEK_SET) != 0)
1559         {
1560                 /*
1561                  * NOTE: Formats that use the lookahead buffer can unset this in their
1562                  * Init routine.
1563                  */
1564                 AH->readHeader = 1;
1565         }
1566         else
1567                 AH->lookaheadLen = 0;   /* Don't bother since we've reset the file */
1568
1569 #if 0
1570         write_msg(modulename, "read %lu bytes into lookahead buffer\n",
1571                           (unsigned long) AH->lookaheadLen);
1572 #endif
1573
1574         /* Close the file */
1575         if (wantClose)
1576                 if (fclose(fh) != 0)
1577                         die_horribly(AH, modulename, "could not close input file: %s\n",
1578                                                  strerror(errno));
1579
1580         return AH->format;
1581 }
1582
1583
1584 /*
1585  * Allocate an archive handle
1586  */
1587 static ArchiveHandle *
1588 _allocAH(const char *FileSpec, const ArchiveFormat fmt,
1589                  const int compression, ArchiveMode mode)
1590 {
1591         ArchiveHandle *AH;
1592
1593 #if 0
1594         write_msg(modulename, "allocating AH for %s, format %d\n", FileSpec, fmt);
1595 #endif
1596
1597         AH = (ArchiveHandle *) calloc(1, sizeof(ArchiveHandle));
1598         if (!AH)
1599                 die_horribly(AH, modulename, "out of memory\n");
1600
1601         /* AH->debugLevel = 100; */
1602
1603         AH->vmaj = K_VERS_MAJOR;
1604         AH->vmin = K_VERS_MINOR;
1605         AH->vrev = K_VERS_REV;
1606
1607         /* initialize for backwards compatible string processing */
1608         AH->public.encoding = PG_SQL_ASCII;
1609         AH->public.std_strings = false;
1610
1611         /* sql error handling */
1612         AH->public.exit_on_error = true;
1613         AH->public.n_errors = 0;
1614
1615         AH->createDate = time(NULL);
1616
1617         AH->intSize = sizeof(int);
1618         AH->offSize = sizeof(off_t);
1619         if (FileSpec)
1620         {
1621                 AH->fSpec = strdup(FileSpec);
1622
1623                 /*
1624                  * Not used; maybe later....
1625                  *
1626                  * AH->workDir = strdup(FileSpec); for(i=strlen(FileSpec) ; i > 0 ;
1627                  * i--) if (AH->workDir[i-1] == '/')
1628                  */
1629         }
1630         else
1631                 AH->fSpec = NULL;
1632
1633         AH->currUser = strdup("");      /* So it's valid, but we can free() it later
1634                                                                  * if necessary */
1635         AH->currSchema = strdup("");    /* ditto */
1636         AH->currWithOids = -1;          /* force SET */
1637
1638         AH->toc = (TocEntry *) calloc(1, sizeof(TocEntry));
1639         if (!AH->toc)
1640                 die_horribly(AH, modulename, "out of memory\n");
1641
1642         AH->toc->next = AH->toc;
1643         AH->toc->prev = AH->toc;
1644
1645         AH->mode = mode;
1646         AH->compression = compression;
1647
1648         AH->pgCopyBuf = createPQExpBuffer();
1649         AH->sqlBuf = createPQExpBuffer();
1650
1651         /* Open stdout with no compression for AH output handle */
1652         AH->gzOut = 0;
1653         AH->OF = stdout;
1654
1655         /*
1656          * On Windows, we need to use binary mode to read/write non-text archive
1657          * formats.  Force stdin/stdout into binary mode if that is what we are
1658          * using.
1659          */
1660 #ifdef WIN32
1661         if (fmt != archNull &&
1662                 (AH->fSpec == NULL || strcmp(AH->fSpec, "") == 0))
1663         {
1664                 if (mode == archModeWrite)
1665                         setmode(fileno(stdout), O_BINARY);
1666                 else
1667                         setmode(fileno(stdin), O_BINARY);
1668         }
1669 #endif
1670
1671 #if 0
1672         write_msg(modulename, "archive format is %d\n", fmt);
1673 #endif
1674
1675         if (fmt == archUnknown)
1676                 AH->format = _discoverArchiveFormat(AH);
1677         else
1678                 AH->format = fmt;
1679
1680         switch (AH->format)
1681         {
1682                 case archCustom:
1683                         InitArchiveFmt_Custom(AH);
1684                         break;
1685
1686                 case archFiles:
1687                         InitArchiveFmt_Files(AH);
1688                         break;
1689
1690                 case archNull:
1691                         InitArchiveFmt_Null(AH);
1692                         break;
1693
1694                 case archTar:
1695                         InitArchiveFmt_Tar(AH);
1696                         break;
1697
1698                 default:
1699                         die_horribly(AH, modulename, "unrecognized file format \"%d\"\n", fmt);
1700         }
1701
1702         return AH;
1703 }
1704
1705
1706 void
1707 WriteDataChunks(ArchiveHandle *AH)
1708 {
1709         TocEntry   *te = AH->toc->next;
1710         StartDataPtr startPtr;
1711         EndDataPtr      endPtr;
1712
1713         while (te != AH->toc)
1714         {
1715                 if (te->dataDumper != NULL)
1716                 {
1717                         AH->currToc = te;
1718                         /* printf("Writing data for %d (%x)\n", te->id, te); */
1719
1720                         if (strcmp(te->desc, "BLOBS") == 0)
1721                         {
1722                                 startPtr = AH->StartBlobsPtr;
1723                                 endPtr = AH->EndBlobsPtr;
1724                         }
1725                         else
1726                         {
1727                                 startPtr = AH->StartDataPtr;
1728                                 endPtr = AH->EndDataPtr;
1729                         }
1730
1731                         if (startPtr != NULL)
1732                                 (*startPtr) (AH, te);
1733
1734                         /*
1735                          * printf("Dumper arg for %d is %x\n", te->id, te->dataDumperArg);
1736                          */
1737
1738                         /*
1739                          * The user-provided DataDumper routine needs to call
1740                          * AH->WriteData
1741                          */
1742                         (*te->dataDumper) ((Archive *) AH, te->dataDumperArg);
1743
1744                         if (endPtr != NULL)
1745                                 (*endPtr) (AH, te);
1746                         AH->currToc = NULL;
1747                 }
1748                 te = te->next;
1749         }
1750 }
1751
1752 void
1753 WriteToc(ArchiveHandle *AH)
1754 {
1755         TocEntry   *te;
1756         char            workbuf[32];
1757         int                     i;
1758
1759         /* printf("%d TOC Entries to save\n", AH->tocCount); */
1760
1761         WriteInt(AH, AH->tocCount);
1762
1763         for (te = AH->toc->next; te != AH->toc; te = te->next)
1764         {
1765                 WriteInt(AH, te->dumpId);
1766                 WriteInt(AH, te->dataDumper ? 1 : 0);
1767
1768                 /* OID is recorded as a string for historical reasons */
1769                 sprintf(workbuf, "%u", te->catalogId.tableoid);
1770                 WriteStr(AH, workbuf);
1771                 sprintf(workbuf, "%u", te->catalogId.oid);
1772                 WriteStr(AH, workbuf);
1773
1774                 WriteStr(AH, te->tag);
1775                 WriteStr(AH, te->desc);
1776                 WriteStr(AH, te->defn);
1777                 WriteStr(AH, te->dropStmt);
1778                 WriteStr(AH, te->copyStmt);
1779                 WriteStr(AH, te->namespace);
1780                 WriteStr(AH, te->tablespace);
1781                 WriteStr(AH, te->owner);
1782                 WriteStr(AH, te->withOids ? "true" : "false");
1783
1784                 /* Dump list of dependencies */
1785                 for (i = 0; i < te->nDeps; i++)
1786                 {
1787                         sprintf(workbuf, "%d", te->dependencies[i]);
1788                         WriteStr(AH, workbuf);
1789                 }
1790                 WriteStr(AH, NULL);             /* Terminate List */
1791
1792                 if (AH->WriteExtraTocPtr)
1793                         (*AH->WriteExtraTocPtr) (AH, te);
1794         }
1795 }
1796
1797 void
1798 ReadToc(ArchiveHandle *AH)
1799 {
1800         int                     i;
1801         char       *tmp;
1802         DumpId     *deps;
1803         int                     depIdx;
1804         int                     depSize;
1805
1806         TocEntry   *te = AH->toc->next;
1807
1808         AH->tocCount = ReadInt(AH);
1809         AH->maxDumpId = 0;
1810
1811         for (i = 0; i < AH->tocCount; i++)
1812         {
1813                 te = (TocEntry *) calloc(1, sizeof(TocEntry));
1814                 te->dumpId = ReadInt(AH);
1815
1816                 if (te->dumpId > AH->maxDumpId)
1817                         AH->maxDumpId = te->dumpId;
1818
1819                 /* Sanity check */
1820                 if (te->dumpId <= 0)
1821                         die_horribly(AH, modulename,
1822                                            "entry ID %d out of range -- perhaps a corrupt TOC\n",
1823                                                  te->dumpId);
1824
1825                 te->hadDumper = ReadInt(AH);
1826
1827                 if (AH->version >= K_VERS_1_8)
1828                 {
1829                         tmp = ReadStr(AH);
1830                         sscanf(tmp, "%u", &te->catalogId.tableoid);
1831                         free(tmp);
1832                 }
1833                 else
1834                         te->catalogId.tableoid = InvalidOid;
1835                 tmp = ReadStr(AH);
1836                 sscanf(tmp, "%u", &te->catalogId.oid);
1837                 free(tmp);
1838
1839                 te->tag = ReadStr(AH);
1840                 te->desc = ReadStr(AH);
1841                 te->defn = ReadStr(AH);
1842                 te->dropStmt = ReadStr(AH);
1843
1844                 if (AH->version >= K_VERS_1_3)
1845                         te->copyStmt = ReadStr(AH);
1846
1847                 if (AH->version >= K_VERS_1_6)
1848                         te->namespace = ReadStr(AH);
1849
1850                 if (AH->version >= K_VERS_1_10)
1851                         te->tablespace = ReadStr(AH);
1852
1853                 te->owner = ReadStr(AH);
1854                 if (AH->version >= K_VERS_1_9)
1855                 {
1856                         if (strcmp(ReadStr(AH), "true") == 0)
1857                                 te->withOids = true;
1858                         else
1859                                 te->withOids = false;
1860                 }
1861                 else
1862                         te->withOids = true;
1863
1864                 /* Read TOC entry dependencies */
1865                 if (AH->version >= K_VERS_1_5)
1866                 {
1867                         depSize = 100;
1868                         deps = (DumpId *) malloc(sizeof(DumpId) * depSize);
1869                         depIdx = 0;
1870                         for (;;)
1871                         {
1872                                 tmp = ReadStr(AH);
1873                                 if (!tmp)
1874                                         break;          /* end of list */
1875                                 if (depIdx >= depSize)
1876                                 {
1877                                         depSize *= 2;
1878                                         deps = (DumpId *) realloc(deps, sizeof(DumpId) * depSize);
1879                                 }
1880                                 sscanf(tmp, "%d", &deps[depIdx]);
1881                                 free(tmp);
1882                                 depIdx++;
1883                         }
1884
1885                         if (depIdx > 0)         /* We have a non-null entry */
1886                         {
1887                                 deps = (DumpId *) realloc(deps, sizeof(DumpId) * depIdx);
1888                                 te->dependencies = deps;
1889                                 te->nDeps = depIdx;
1890                         }
1891                         else
1892                         {
1893                                 free(deps);
1894                                 te->dependencies = NULL;
1895                                 te->nDeps = 0;
1896                         }
1897                 }
1898                 else
1899                 {
1900                         te->dependencies = NULL;
1901                         te->nDeps = 0;
1902                 }
1903
1904                 if (AH->ReadExtraTocPtr)
1905                         (*AH->ReadExtraTocPtr) (AH, te);
1906
1907                 ahlog(AH, 3, "read TOC entry %d (ID %d) for %s %s\n",
1908                           i, te->dumpId, te->desc, te->tag);
1909
1910                 /* link completed entry into TOC circular list */
1911                 te->prev = AH->toc->prev;
1912                 AH->toc->prev->next = te;
1913                 AH->toc->prev = te;
1914                 te->next = AH->toc;
1915
1916                 /* special processing immediately upon read for some items */
1917                 if (strcmp(te->desc, "ENCODING") == 0)
1918                         processEncodingEntry(AH, te);
1919                 else if (strcmp(te->desc, "STDSTRINGS") == 0)
1920                         processStdStringsEntry(AH, te);
1921         }
1922 }
1923
1924 static void
1925 processEncodingEntry(ArchiveHandle *AH, TocEntry *te)
1926 {
1927         /* te->defn should have the form SET client_encoding = 'foo'; */
1928         char       *defn = strdup(te->defn);
1929         char       *ptr1;
1930         char       *ptr2 = NULL;
1931         int                     encoding;
1932
1933         ptr1 = strchr(defn, '\'');
1934         if (ptr1)
1935                 ptr2 = strchr(++ptr1, '\'');
1936         if (ptr2)
1937         {
1938                 *ptr2 = '\0';
1939                 encoding = pg_char_to_encoding(ptr1);
1940                 if (encoding < 0)
1941                         die_horribly(AH, modulename, "unrecognized encoding \"%s\"\n",
1942                                                  ptr1);
1943                 AH->public.encoding = encoding;
1944         }
1945         else
1946                 die_horribly(AH, modulename, "invalid ENCODING item: %s\n",
1947                                          te->defn);
1948
1949         free(defn);
1950 }
1951
1952 static void
1953 processStdStringsEntry(ArchiveHandle *AH, TocEntry *te)
1954 {
1955         /* te->defn should have the form SET standard_conforming_strings = 'x'; */
1956         char       *ptr1;
1957
1958         ptr1 = strchr(te->defn, '\'');
1959         if (ptr1 && strncmp(ptr1, "'on'", 4) == 0)
1960                 AH->public.std_strings = true;
1961         else if (ptr1 && strncmp(ptr1, "'off'", 5) == 0)
1962                 AH->public.std_strings = false;
1963         else
1964                 die_horribly(AH, modulename, "invalid STDSTRINGS item: %s\n",
1965                                          te->defn);
1966 }
1967
1968 static teReqs
1969 _tocEntryRequired(TocEntry *te, RestoreOptions *ropt, bool include_acls)
1970 {
1971         teReqs          res = REQ_ALL;
1972
1973         /* ENCODING and STDSTRINGS items are dumped specially, so always reject */
1974         if (strcmp(te->desc, "ENCODING") == 0 ||
1975                 strcmp(te->desc, "STDSTRINGS") == 0)
1976                 return 0;
1977
1978         /* If it's an ACL, maybe ignore it */
1979         if ((!include_acls || ropt->aclsSkip) && strcmp(te->desc, "ACL") == 0)
1980                 return 0;
1981
1982         if (!ropt->create && strcmp(te->desc, "DATABASE") == 0)
1983                 return 0;
1984
1985         /* Check options for selective dump/restore */
1986         if (ropt->schemaNames)
1987         {
1988                 /* If no namespace is specified, it means all. */
1989                 if (!te->namespace)
1990                         return 0;
1991                 if (strcmp(ropt->schemaNames, te->namespace) != 0)
1992                         return 0;
1993         }
1994
1995         if (ropt->selTypes)
1996         {
1997                 if (strcmp(te->desc, "TABLE") == 0 ||
1998                         strcmp(te->desc, "TABLE DATA") == 0)
1999                 {
2000                         if (!ropt->selTable)
2001                                 return 0;
2002                         if (ropt->tableNames && strcmp(ropt->tableNames, te->tag) != 0)
2003                                 return 0;
2004                 }
2005                 else if (strcmp(te->desc, "INDEX") == 0)
2006                 {
2007                         if (!ropt->selIndex)
2008                                 return 0;
2009                         if (ropt->indexNames && strcmp(ropt->indexNames, te->tag) != 0)
2010                                 return 0;
2011                 }
2012                 else if (strcmp(te->desc, "FUNCTION") == 0)
2013                 {
2014                         if (!ropt->selFunction)
2015                                 return 0;
2016                         if (ropt->functionNames && strcmp(ropt->functionNames, te->tag) != 0)
2017                                 return 0;
2018                 }
2019                 else if (strcmp(te->desc, "TRIGGER") == 0)
2020                 {
2021                         if (!ropt->selTrigger)
2022                                 return 0;
2023                         if (ropt->triggerNames && strcmp(ropt->triggerNames, te->tag) != 0)
2024                                 return 0;
2025                 }
2026                 else
2027                         return 0;
2028         }
2029
2030         /*
2031          * Check if we had a dataDumper. Indicates if the entry is schema or data
2032          */
2033         if (!te->hadDumper)
2034         {
2035                 /*
2036                  * Special Case: If 'SEQUENCE SET' then it is considered a data entry
2037                  */
2038                 if (strcmp(te->desc, "SEQUENCE SET") == 0)
2039                         res = res & REQ_DATA;
2040                 else
2041                         res = res & ~REQ_DATA;
2042         }
2043
2044         /*
2045          * Special case: <Init> type with <Max OID> tag; this is obsolete and we
2046          * always ignore it.
2047          */
2048         if ((strcmp(te->desc, "<Init>") == 0) && (strcmp(te->tag, "Max OID") == 0))
2049                 return 0;
2050
2051         /* Mask it if we only want schema */
2052         if (ropt->schemaOnly)
2053                 res = res & REQ_SCHEMA;
2054
2055         /* Mask it we only want data */
2056         if (ropt->dataOnly)
2057                 res = res & REQ_DATA;
2058
2059         /* Mask it if we don't have a schema contribution */
2060         if (!te->defn || strlen(te->defn) == 0)
2061                 res = res & ~REQ_SCHEMA;
2062
2063         /* Finally, if we used a list, limit based on that as well */
2064         if (ropt->limitToList && !ropt->idWanted[te->dumpId - 1])
2065                 return 0;
2066
2067         return res;
2068 }
2069
2070 /*
2071  * Issue SET commands for parameters that we want to have set the same way
2072  * at all times during execution of a restore script.
2073  */
2074 static void
2075 _doSetFixedOutputState(ArchiveHandle *AH)
2076 {
2077         /* Select the correct character set encoding */
2078         ahprintf(AH, "SET client_encoding = '%s';\n",
2079                          pg_encoding_to_char(AH->public.encoding));
2080
2081         /* Select the correct string literal syntax */
2082         ahprintf(AH, "SET standard_conforming_strings = %s;\n",
2083                          AH->public.std_strings ? "on" : "off");
2084
2085         /* Make sure function checking is disabled */
2086         ahprintf(AH, "SET check_function_bodies = false;\n");
2087
2088         /* Avoid annoying notices etc */
2089         ahprintf(AH, "SET client_min_messages = warning;\n");
2090         if (!AH->public.std_strings)
2091                 ahprintf(AH, "SET escape_string_warning = off;\n");
2092
2093         ahprintf(AH, "\n");
2094 }
2095
2096 /*
2097  * Issue a SET SESSION AUTHORIZATION command.  Caller is responsible
2098  * for updating state if appropriate.  If user is NULL or an empty string,
2099  * the specification DEFAULT will be used.
2100  */
2101 static void
2102 _doSetSessionAuth(ArchiveHandle *AH, const char *user)
2103 {
2104         PQExpBuffer cmd = createPQExpBuffer();
2105
2106         appendPQExpBuffer(cmd, "SET SESSION AUTHORIZATION ");
2107
2108         /*
2109          * SQL requires a string literal here.  Might as well be correct.
2110          */
2111         if (user && *user)
2112                 appendStringLiteralAHX(cmd, user, AH);
2113         else
2114                 appendPQExpBuffer(cmd, "DEFAULT");
2115         appendPQExpBuffer(cmd, ";");
2116
2117         if (RestoringToDB(AH))
2118         {
2119                 PGresult   *res;
2120
2121                 res = PQexec(AH->connection, cmd->data);
2122
2123                 if (!res || PQresultStatus(res) != PGRES_COMMAND_OK)
2124                         /* NOT warn_or_die_horribly... use -O instead to skip this. */
2125                         die_horribly(AH, modulename, "could not set session user to \"%s\": %s",
2126                                                  user, PQerrorMessage(AH->connection));
2127
2128                 PQclear(res);
2129         }
2130         else
2131                 ahprintf(AH, "%s\n\n", cmd->data);
2132
2133         destroyPQExpBuffer(cmd);
2134 }
2135
2136
2137 /*
2138  * Issue a SET default_with_oids command.  Caller is responsible
2139  * for updating state if appropriate.
2140  */
2141 static void
2142 _doSetWithOids(ArchiveHandle *AH, const bool withOids)
2143 {
2144         PQExpBuffer cmd = createPQExpBuffer();
2145
2146         appendPQExpBuffer(cmd, "SET default_with_oids = %s;", withOids ?
2147                                           "true" : "false");
2148
2149         if (RestoringToDB(AH))
2150         {
2151                 PGresult   *res;
2152
2153                 res = PQexec(AH->connection, cmd->data);
2154
2155                 if (!res || PQresultStatus(res) != PGRES_COMMAND_OK)
2156                         warn_or_die_horribly(AH, modulename,
2157                                                                  "could not set default_with_oids: %s",
2158                                                                  PQerrorMessage(AH->connection));
2159
2160                 PQclear(res);
2161         }
2162         else
2163                 ahprintf(AH, "%s\n\n", cmd->data);
2164
2165         destroyPQExpBuffer(cmd);
2166 }
2167
2168
2169 /*
2170  * Issue the commands to connect to the specified database.
2171  *
2172  * If we're currently restoring right into a database, this will
2173  * actually establish a connection. Otherwise it puts a \connect into
2174  * the script output.
2175  *
2176  * NULL dbname implies reconnecting to the current DB (pretty useless).
2177  */
2178 static void
2179 _reconnectToDB(ArchiveHandle *AH, const char *dbname)
2180 {
2181         if (RestoringToDB(AH))
2182                 ReconnectToServer(AH, dbname, NULL);
2183         else
2184         {
2185                 PQExpBuffer qry = createPQExpBuffer();
2186
2187                 appendPQExpBuffer(qry, "\\connect %s\n\n",
2188                                                   dbname ? fmtId(dbname) : "-");
2189                 ahprintf(AH, "%s", qry->data);
2190                 destroyPQExpBuffer(qry);
2191         }
2192
2193         /*
2194          * NOTE: currUser keeps track of what the imaginary session user in our
2195          * script is.  It's now effectively reset to the original userID.
2196          */
2197         if (AH->currUser)
2198                 free(AH->currUser);
2199
2200         AH->currUser = strdup("");
2201
2202         /* don't assume we still know the output schema */
2203         if (AH->currSchema)
2204                 free(AH->currSchema);
2205         AH->currSchema = strdup("");
2206         AH->currWithOids = -1;
2207
2208         /* re-establish fixed state */
2209         _doSetFixedOutputState(AH);
2210 }
2211
2212 /*
2213  * Become the specified user, and update state to avoid redundant commands
2214  *
2215  * NULL or empty argument is taken to mean restoring the session default
2216  */
2217 static void
2218 _becomeUser(ArchiveHandle *AH, const char *user)
2219 {
2220         if (!user)
2221                 user = "";                              /* avoid null pointers */
2222
2223         if (AH->currUser && strcmp(AH->currUser, user) == 0)
2224                 return;                                 /* no need to do anything */
2225
2226         _doSetSessionAuth(AH, user);
2227
2228         /*
2229          * NOTE: currUser keeps track of what the imaginary session user in our
2230          * script is
2231          */
2232         if (AH->currUser)
2233                 free(AH->currUser);
2234
2235         AH->currUser = strdup(user);
2236 }
2237
2238 /*
2239  * Become the owner of the the given TOC entry object.  If
2240  * changes in ownership are not allowed, this doesn't do anything.
2241  */
2242 static void
2243 _becomeOwner(ArchiveHandle *AH, TocEntry *te)
2244 {
2245         if (AH->ropt && (AH->ropt->noOwner || !AH->ropt->use_setsessauth))
2246                 return;
2247
2248         _becomeUser(AH, te->owner);
2249 }
2250
2251
2252 /*
2253  * Set the proper default_with_oids value for the table.
2254  */
2255 static void
2256 _setWithOids(ArchiveHandle *AH, TocEntry *te)
2257 {
2258         if (AH->currWithOids != te->withOids)
2259         {
2260                 _doSetWithOids(AH, te->withOids);
2261                 AH->currWithOids = te->withOids;
2262         }
2263 }
2264
2265
2266 /*
2267  * Issue the commands to select the specified schema as the current schema
2268  * in the target database.
2269  */
2270 static void
2271 _selectOutputSchema(ArchiveHandle *AH, const char *schemaName)
2272 {
2273         PQExpBuffer qry;
2274
2275         if (!schemaName || *schemaName == '\0' ||
2276                 (AH->currSchema && strcmp(AH->currSchema, schemaName) == 0))
2277                 return;                                 /* no need to do anything */
2278
2279         qry = createPQExpBuffer();
2280
2281         appendPQExpBuffer(qry, "SET search_path = %s",
2282                                           fmtId(schemaName));
2283         if (strcmp(schemaName, "pg_catalog") != 0)
2284                 appendPQExpBuffer(qry, ", pg_catalog");
2285
2286         if (RestoringToDB(AH))
2287         {
2288                 PGresult   *res;
2289
2290                 res = PQexec(AH->connection, qry->data);
2291
2292                 if (!res || PQresultStatus(res) != PGRES_COMMAND_OK)
2293                         warn_or_die_horribly(AH, modulename,
2294                                                                  "could not set search_path to \"%s\": %s",
2295                                                                  schemaName, PQerrorMessage(AH->connection));
2296
2297                 PQclear(res);
2298         }
2299         else
2300                 ahprintf(AH, "%s;\n\n", qry->data);
2301
2302         if (AH->currSchema)
2303                 free(AH->currSchema);
2304         AH->currSchema = strdup(schemaName);
2305
2306         destroyPQExpBuffer(qry);
2307 }
2308
2309 /*
2310  * Issue the commands to select the specified tablespace as the current one
2311  * in the target database.
2312  */
2313 static void
2314 _selectTablespace(ArchiveHandle *AH, const char *tablespace)
2315 {
2316         PQExpBuffer qry;
2317         const char *want,
2318                            *have;
2319
2320         have = AH->currTablespace;
2321         want = tablespace;
2322
2323         /* no need to do anything for non-tablespace object */
2324         if (!want)
2325                 return;
2326
2327         if (have && strcmp(want, have) == 0)
2328                 return;                                 /* no need to do anything */
2329
2330         qry = createPQExpBuffer();
2331
2332         if (strcmp(want, "") == 0)
2333         {
2334                 /* We want the tablespace to be the database's default */
2335                 appendPQExpBuffer(qry, "SET default_tablespace = ''");
2336         }
2337         else
2338         {
2339                 /* We want an explicit tablespace */
2340                 appendPQExpBuffer(qry, "SET default_tablespace = %s", fmtId(want));
2341         }
2342
2343         if (RestoringToDB(AH))
2344         {
2345                 PGresult   *res;
2346
2347                 res = PQexec(AH->connection, qry->data);
2348
2349                 if (!res || PQresultStatus(res) != PGRES_COMMAND_OK)
2350                         warn_or_die_horribly(AH, modulename,
2351                                                                  "could not set default_tablespace to %s: %s",
2352                                                                  fmtId(want), PQerrorMessage(AH->connection));
2353
2354                 PQclear(res);
2355         }
2356         else
2357                 ahprintf(AH, "%s;\n\n", qry->data);
2358
2359         if (AH->currTablespace)
2360                 free(AH->currTablespace);
2361         AH->currTablespace = strdup(want);
2362
2363         destroyPQExpBuffer(qry);
2364 }
2365
2366 /*
2367  * Extract an object description for a TOC entry, and append it to buf.
2368  *
2369  * This is not quite as general as it may seem, since it really only
2370  * handles constructing the right thing to put into ALTER ... OWNER TO.
2371  *
2372  * The whole thing is pretty grotty, but we are kind of stuck since the
2373  * information used is all that's available in older dump files.
2374  */
2375 static void
2376 _getObjectDescription(PQExpBuffer buf, TocEntry *te, ArchiveHandle *AH)
2377 {
2378         const char *type = te->desc;
2379
2380         /* Use ALTER TABLE for views and sequences */
2381         if (strcmp(type, "VIEW") == 0 || strcmp(type, "SEQUENCE") == 0)
2382                 type = "TABLE";
2383
2384         /* objects named by a schema and name */
2385         if (strcmp(type, "CONVERSION") == 0 ||
2386                 strcmp(type, "DOMAIN") == 0 ||
2387                 strcmp(type, "TABLE") == 0 ||
2388                 strcmp(type, "TYPE") == 0)
2389         {
2390                 appendPQExpBuffer(buf, "%s ", type);
2391                 if (te->namespace && te->namespace[0])  /* is null pre-7.3 */
2392                         appendPQExpBuffer(buf, "%s.", fmtId(te->namespace));
2393
2394                 /*
2395                  * Pre-7.3 pg_dump would sometimes (not always) put a fmtId'd name
2396                  * into te->tag for an index. This check is heuristic, so make its
2397                  * scope as narrow as possible.
2398                  */
2399                 if (AH->version < K_VERS_1_7 &&
2400                         te->tag[0] == '"' &&
2401                         te->tag[strlen(te->tag) - 1] == '"' &&
2402                         strcmp(type, "INDEX") == 0)
2403                         appendPQExpBuffer(buf, "%s", te->tag);
2404                 else
2405                         appendPQExpBuffer(buf, "%s", fmtId(te->tag));
2406                 return;
2407         }
2408
2409         /* objects named by just a name */
2410         if (strcmp(type, "DATABASE") == 0 ||
2411                 strcmp(type, "SCHEMA") == 0)
2412         {
2413                 appendPQExpBuffer(buf, "%s %s", type, fmtId(te->tag));
2414                 return;
2415         }
2416
2417         /*
2418          * These object types require additional decoration.  Fortunately, the
2419          * information needed is exactly what's in the DROP command.
2420          */
2421         if (strcmp(type, "AGGREGATE") == 0 ||
2422                 strcmp(type, "FUNCTION") == 0 ||
2423                 strcmp(type, "OPERATOR") == 0 ||
2424                 strcmp(type, "OPERATOR CLASS") == 0)
2425         {
2426                 /* Chop "DROP " off the front and make a modifiable copy */
2427                 char       *first = strdup(te->dropStmt + 5);
2428                 char       *last;
2429
2430                 /* point to last character in string */
2431                 last = first + strlen(first) - 1;
2432
2433                 /* Strip off any ';' or '\n' at the end */
2434                 while (last >= first && (*last == '\n' || *last == ';'))
2435                         last--;
2436                 *(last + 1) = '\0';
2437
2438                 appendPQExpBufferStr(buf, first);
2439
2440                 free(first);
2441                 return;
2442         }
2443
2444         write_msg(modulename, "WARNING: don't know how to set owner for object type %s\n",
2445                           type);
2446 }
2447
2448 static void
2449 _printTocEntry(ArchiveHandle *AH, TocEntry *te, RestoreOptions *ropt, bool isData, bool acl_pass)
2450 {
2451         /* ACLs are dumped only during acl pass */
2452         if (acl_pass)
2453         {
2454                 if (strcmp(te->desc, "ACL") != 0)
2455                         return;
2456         }
2457         else
2458         {
2459                 if (strcmp(te->desc, "ACL") == 0)
2460                         return;
2461         }
2462
2463         /*
2464          * Avoid dumping the public schema, as it will already be created ...
2465          * unless we are using --clean mode, in which case it's been deleted and
2466          * we'd better recreate it.
2467          */
2468         if (!ropt->dropSchema &&
2469                 strcmp(te->desc, "SCHEMA") == 0 && strcmp(te->tag, "public") == 0)
2470                 return;
2471
2472         /* Select owner, schema, and tablespace as necessary */
2473         _becomeOwner(AH, te);
2474         _selectOutputSchema(AH, te->namespace);
2475         _selectTablespace(AH, te->tablespace);
2476
2477         /* Set up OID mode too */
2478         if (strcmp(te->desc, "TABLE") == 0)
2479                 _setWithOids(AH, te);
2480
2481         /* Emit header comment for item */
2482         if (!AH->noTocComments)
2483         {
2484                 const char *pfx;
2485
2486                 if (isData)
2487                         pfx = "Data for ";
2488                 else
2489                         pfx = "";
2490
2491                 ahprintf(AH, "--\n");
2492                 if (AH->public.verbose)
2493                 {
2494                         ahprintf(AH, "-- TOC entry %d (class %u OID %u)\n",
2495                                          te->dumpId, te->catalogId.tableoid, te->catalogId.oid);
2496                         if (te->nDeps > 0)
2497                         {
2498                                 int                     i;
2499
2500                                 ahprintf(AH, "-- Dependencies:");
2501                                 for (i = 0; i < te->nDeps; i++)
2502                                         ahprintf(AH, " %d", te->dependencies[i]);
2503                                 ahprintf(AH, "\n");
2504                         }
2505                 }
2506                 ahprintf(AH, "-- %sName: %s; Type: %s; Schema: %s; Owner: %s",
2507                                  pfx, te->tag, te->desc,
2508                                  te->namespace ? te->namespace : "-",
2509                                  ropt->noOwner ? "-" : te->owner);
2510                 if (te->tablespace)
2511                         ahprintf(AH, "; Tablespace: %s", te->tablespace);
2512                 ahprintf(AH, "\n");
2513
2514                 if (AH->PrintExtraTocPtr !=NULL)
2515                         (*AH->PrintExtraTocPtr) (AH, te);
2516                 ahprintf(AH, "--\n\n");
2517         }
2518
2519         /*
2520          * Actually print the definition.
2521          *
2522          * Really crude hack for suppressing AUTHORIZATION clause that old pg_dump
2523          * versions put into CREATE SCHEMA.  We have to do this when --no-owner
2524          * mode is selected.  This is ugly, but I see no other good way ...
2525          */
2526         if (ropt->noOwner && strcmp(te->desc, "SCHEMA") == 0)
2527         {
2528                 ahprintf(AH, "CREATE SCHEMA %s;\n\n\n", fmtId(te->tag));
2529         }
2530         else
2531         {
2532                 if (strlen(te->defn) > 0)
2533                         ahprintf(AH, "%s\n\n", te->defn);
2534         }
2535
2536         /*
2537          * If we aren't using SET SESSION AUTH to determine ownership, we must
2538          * instead issue an ALTER OWNER command.  We assume that anything without
2539          * a DROP command is not a separately ownable object.  All the categories
2540          * with DROP commands must appear in one list or the other.
2541          */
2542         if (!ropt->noOwner && !ropt->use_setsessauth &&
2543                 strlen(te->owner) > 0 && strlen(te->dropStmt) > 0)
2544         {
2545                 if (strcmp(te->desc, "AGGREGATE") == 0 ||
2546                         strcmp(te->desc, "CONVERSION") == 0 ||
2547                         strcmp(te->desc, "DATABASE") == 0 ||
2548                         strcmp(te->desc, "DOMAIN") == 0 ||
2549                         strcmp(te->desc, "FUNCTION") == 0 ||
2550                         strcmp(te->desc, "OPERATOR") == 0 ||
2551                         strcmp(te->desc, "OPERATOR CLASS") == 0 ||
2552                         strcmp(te->desc, "SCHEMA") == 0 ||
2553                         strcmp(te->desc, "TABLE") == 0 ||
2554                         strcmp(te->desc, "TYPE") == 0 ||
2555                         strcmp(te->desc, "VIEW") == 0 ||
2556                         strcmp(te->desc, "SEQUENCE") == 0)
2557                 {
2558                         PQExpBuffer temp = createPQExpBuffer();
2559
2560                         appendPQExpBuffer(temp, "ALTER ");
2561                         _getObjectDescription(temp, te, AH);
2562                         appendPQExpBuffer(temp, " OWNER TO %s;", fmtId(te->owner));
2563                         ahprintf(AH, "%s\n\n", temp->data);
2564                         destroyPQExpBuffer(temp);
2565                 }
2566                 else if (strcmp(te->desc, "CAST") == 0 ||
2567                                  strcmp(te->desc, "CHECK CONSTRAINT") == 0 ||
2568                                  strcmp(te->desc, "CONSTRAINT") == 0 ||
2569                                  strcmp(te->desc, "DEFAULT") == 0 ||
2570                                  strcmp(te->desc, "FK CONSTRAINT") == 0 ||
2571                                  strcmp(te->desc, "INDEX") == 0 ||
2572                                  strcmp(te->desc, "PROCEDURAL LANGUAGE") == 0 ||
2573                                  strcmp(te->desc, "RULE") == 0 ||
2574                                  strcmp(te->desc, "TRIGGER") == 0)
2575                 {
2576                         /* these object types don't have separate owners */
2577                 }
2578                 else
2579                 {
2580                         write_msg(modulename, "WARNING: don't know how to set owner for object type %s\n",
2581                                           te->desc);
2582                 }
2583         }
2584
2585         /*
2586          * If it's an ACL entry, it might contain SET SESSION AUTHORIZATION
2587          * commands, so we can no longer assume we know the current auth setting.
2588          */
2589         if (strncmp(te->desc, "ACL", 3) == 0)
2590         {
2591                 if (AH->currUser)
2592                         free(AH->currUser);
2593                 AH->currUser = NULL;
2594         }
2595 }
2596
2597 void
2598 WriteHead(ArchiveHandle *AH)
2599 {
2600         struct tm       crtm;
2601
2602         (*AH->WriteBufPtr) (AH, "PGDMP", 5);            /* Magic code */
2603         (*AH->WriteBytePtr) (AH, AH->vmaj);
2604         (*AH->WriteBytePtr) (AH, AH->vmin);
2605         (*AH->WriteBytePtr) (AH, AH->vrev);
2606         (*AH->WriteBytePtr) (AH, AH->intSize);
2607         (*AH->WriteBytePtr) (AH, AH->offSize);
2608         (*AH->WriteBytePtr) (AH, AH->format);
2609
2610 #ifndef HAVE_LIBZ
2611         if (AH->compression != 0)
2612                 write_msg(modulename, "WARNING: requested compression not available in this "
2613                                   "installation -- archive will be uncompressed\n");
2614
2615         AH->compression = 0;
2616 #endif
2617
2618         WriteInt(AH, AH->compression);
2619
2620         crtm = *localtime(&AH->createDate);
2621         WriteInt(AH, crtm.tm_sec);
2622         WriteInt(AH, crtm.tm_min);
2623         WriteInt(AH, crtm.tm_hour);
2624         WriteInt(AH, crtm.tm_mday);
2625         WriteInt(AH, crtm.tm_mon);
2626         WriteInt(AH, crtm.tm_year);
2627         WriteInt(AH, crtm.tm_isdst);
2628         WriteStr(AH, PQdb(AH->connection));
2629         WriteStr(AH, AH->public.remoteVersionStr);
2630         WriteStr(AH, PG_VERSION);
2631 }
2632
2633 void
2634 ReadHead(ArchiveHandle *AH)
2635 {
2636         char            tmpMag[7];
2637         int                     fmt;
2638         struct tm       crtm;
2639
2640         /* If we haven't already read the header... */
2641         if (!AH->readHeader)
2642         {
2643
2644                 (*AH->ReadBufPtr) (AH, tmpMag, 5);
2645
2646                 if (strncmp(tmpMag, "PGDMP", 5) != 0)
2647                         die_horribly(AH, modulename, "did not find magic string in file header\n");
2648
2649                 AH->vmaj = (*AH->ReadBytePtr) (AH);
2650                 AH->vmin = (*AH->ReadBytePtr) (AH);
2651
2652                 if (AH->vmaj > 1 || ((AH->vmaj == 1) && (AH->vmin > 0)))                /* Version > 1.0 */
2653                         AH->vrev = (*AH->ReadBytePtr) (AH);
2654                 else
2655                         AH->vrev = 0;
2656
2657                 AH->version = ((AH->vmaj * 256 + AH->vmin) * 256 + AH->vrev) * 256 + 0;
2658
2659
2660                 if (AH->version < K_VERS_1_0 || AH->version > K_VERS_MAX)
2661                         die_horribly(AH, modulename, "unsupported version (%d.%d) in file header\n",
2662                                                  AH->vmaj, AH->vmin);
2663
2664                 AH->intSize = (*AH->ReadBytePtr) (AH);
2665                 if (AH->intSize > 32)
2666                         die_horribly(AH, modulename, "sanity check on integer size (%lu) failed\n",
2667                                                  (unsigned long) AH->intSize);
2668
2669                 if (AH->intSize > sizeof(int))
2670                         write_msg(modulename, "WARNING: archive was made on a machine with larger integers, some operations may fail\n");
2671
2672                 if (AH->version >= K_VERS_1_7)
2673                         AH->offSize = (*AH->ReadBytePtr) (AH);
2674                 else
2675                         AH->offSize = AH->intSize;
2676
2677                 fmt = (*AH->ReadBytePtr) (AH);
2678
2679                 if (AH->format != fmt)
2680                         die_horribly(AH, modulename, "expected format (%d) differs from format found in file (%d)\n",
2681                                                  AH->format, fmt);
2682         }
2683
2684         if (AH->version >= K_VERS_1_2)
2685         {
2686                 if (AH->version < K_VERS_1_4)
2687                         AH->compression = (*AH->ReadBytePtr) (AH);
2688                 else
2689                         AH->compression = ReadInt(AH);
2690         }
2691         else
2692                 AH->compression = Z_DEFAULT_COMPRESSION;
2693
2694 #ifndef HAVE_LIBZ
2695         if (AH->compression != 0)
2696                 write_msg(modulename, "WARNING: archive is compressed, but this installation does not support compression -- no data will be available\n");
2697 #endif
2698
2699         if (AH->version >= K_VERS_1_4)
2700         {
2701                 crtm.tm_sec = ReadInt(AH);
2702                 crtm.tm_min = ReadInt(AH);
2703                 crtm.tm_hour = ReadInt(AH);
2704                 crtm.tm_mday = ReadInt(AH);
2705                 crtm.tm_mon = ReadInt(AH);
2706                 crtm.tm_year = ReadInt(AH);
2707                 crtm.tm_isdst = ReadInt(AH);
2708
2709                 AH->archdbname = ReadStr(AH);
2710
2711                 AH->createDate = mktime(&crtm);
2712
2713                 if (AH->createDate == (time_t) -1)
2714                         write_msg(modulename, "WARNING: invalid creation date in header\n");
2715         }
2716
2717         if (AH->version >= K_VERS_1_10)
2718         {
2719                 AH->archiveRemoteVersion = ReadStr(AH);
2720                 AH->archiveDumpVersion = ReadStr(AH);
2721         }
2722
2723 }
2724
2725
2726 /*
2727  * checkSeek
2728  *        check to see if fseek can be performed.
2729  */
2730
2731 bool
2732 checkSeek(FILE *fp)
2733 {
2734
2735         if (fseeko(fp, 0, SEEK_CUR) != 0)
2736                 return false;
2737         else if (sizeof(off_t) > sizeof(long))
2738
2739                 /*
2740                  * At this point, off_t is too large for long, so we return based on
2741                  * whether an off_t version of fseek is available.
2742                  */
2743 #ifdef HAVE_FSEEKO
2744                 return true;
2745 #else
2746                 return false;
2747 #endif
2748         else
2749                 return true;
2750 }
2751
2752
2753 /*
2754  * dumpTimestamp
2755  */
2756 static void
2757 dumpTimestamp(ArchiveHandle *AH, const char *msg, time_t tim)
2758 {
2759         char            buf[256];
2760
2761         if (strftime(buf, 256, "%Y-%m-%d %H:%M:%S %Z", localtime(&tim)) != 0)
2762                 ahprintf(AH, "-- %s %s\n\n", msg, buf);
2763 }