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