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