truncate_restart.diff

diff --git a/doc/src/sgml/ref/truncate.sgml b/doc/src/sgml/ref/truncate.sgml
index f32d255..137eade 100644
*** a/doc/src/sgml/ref/truncate.sgml
--- b/doc/src/sgml/ref/truncate.sgml
*************** TRUNCATE [ TABLE ] [ ONLY ] <replaceable
*** 159,190 ****
     transaction does not commit.
    </para>
  
-   <warning>
-    <para>
-     Any <command>ALTER SEQUENCE RESTART</> operations performed as a
-     consequence of using the <literal>RESTART IDENTITY</> option are
-     nontransactional and will not be rolled back on failure.  To minimize
-     the risk, these operations are performed only after all the rest of
-     <command>TRUNCATE</>'s work is done.  However, there is still a risk
-     if <command>TRUNCATE</> is performed inside a transaction block that is
-     aborted afterwards.  For example, consider
  
- <programlisting>
- BEGIN;
- TRUNCATE TABLE foo RESTART IDENTITY;
- COPY foo FROM ...;
- COMMIT;
- </programlisting>
- 
-     If the <command>COPY</> fails partway through, the table data
-     rolls back correctly, but the sequences will be left with values
-     that are probably smaller than they had before, possibly leading
-     to duplicate-key failures or other problems in later transactions.
-     If this is likely to be a problem, it's best to avoid using
-     <literal>RESTART IDENTITY</>, and accept that the new contents of
-     the table will have higher serial numbers than the old.
-    </para>
-   </warning>
   </refsect1>
  
   <refsect1>
--- 159,165 ----
diff --git a/src/backend/commands/sequence.c b/src/backend/commands/sequence.c
index 04b0c71..4fb9093 100644
*** a/src/backend/commands/sequence.c
--- b/src/backend/commands/sequence.c
***************
*** 35,41 ****
  #include "utils/lsyscache.h"
  #include "utils/resowner.h"
  #include "utils/syscache.h"
! 
  
  /*
   * We don't want to log each fetching of a value from a sequence,
--- 35,41 ----
  #include "utils/lsyscache.h"
  #include "utils/resowner.h"
  #include "utils/syscache.h"
! #include "utils/snapmgr.h"
  
  /*
   * We don't want to log each fetching of a value from a sequence,
*************** static void init_params(List *options, b
*** 96,101 ****
--- 96,104 ----
  static void do_setval(Oid relid, int64 next, bool iscalled);
  static void process_owned_by(Relation seqrel, List *owned_by);
  
+ static void init_seq_relation(Relation rel,TupleDesc tupDesc,Datum * value,
+ 							  bool * null,List * owned_by);
+ 
  
  /*
   * DefineSequence
*************** DefineSequence(CreateSeqStmt *seq)
*** 109,119 ****
  	CreateStmt *stmt = makeNode(CreateStmt);
  	Oid			seqoid;
  	Relation	rel;
! 	Buffer		buf;
! 	Page		page;
! 	sequence_magic *sm;
! 	HeapTuple	tuple;
! 	TupleDesc	tupDesc;
  	Datum		value[SEQ_COL_LASTCOL];
  	bool		null[SEQ_COL_LASTCOL];
  	int			i;
--- 112,118 ----
  	CreateStmt *stmt = makeNode(CreateStmt);
  	Oid			seqoid;
  	Relation	rel;
! 	TupleDesc tupDesc;
  	Datum		value[SEQ_COL_LASTCOL];
  	bool		null[SEQ_COL_LASTCOL];
  	int			i;
*************** DefineSequence(CreateSeqStmt *seq)
*** 210,217 ****
  
  	rel = heap_open(seqoid, AccessExclusiveLock);
  	tupDesc = RelationGetDescr(rel);
  
! 	/* Initialize first page of relation with special magic number */
  
  	buf = ReadBuffer(rel, P_NEW);
  	Assert(BufferGetBlockNumber(buf) == 0);
--- 209,293 ----
  
  	rel = heap_open(seqoid, AccessExclusiveLock);
  	tupDesc = RelationGetDescr(rel);
+ 	init_seq_relation(rel, tupDesc,value,null,owned_by);
+ 	heap_close(rel, NoLock);
+ }
  
! /**
!  * Resets the relation used by a sequence.
!  *
!  * The sequence is reset to its initial values,
!  * the old sequence is left in place in case the
!  * transaction rolls back.
!  */
! void ResetSequenceRelation(Oid seq_relid,List * options)
! {
! 	Relation seq_rel = relation_open(seq_relid,AccessExclusiveLock);
! 	SeqTable elm;
! 	Page page;
! 	Form_pg_sequence seq;
! 	Buffer buf;
! 	TupleDesc tupDesc;
! 	sequence_magic * sm;
! 	HeapTupleData tuple;
! 	ItemId lp;
! 	Datum * values;
! 	bool * isnull;
! 
! 	/**
! 	 * Read the old sequence.
! 	 *
! 	 */
! 	init_sequence(seq_relid,&elm,&seq_rel);
! 	seq = read_info(elm,seq_rel,&buf);
! 	page = BufferGetPage(buf);
! 	sm = (sequence_magic *) PageGetSpecialPointer(page);
! 
! 	if (sm->magic != SEQ_MAGIC)
! 		elog(ERROR, "bad magic number in sequence \"%s\": %08X",
! 			 RelationGetRelationName(seq_rel), sm->magic);
! 
! 	lp = PageGetItemId(page, FirstOffsetNumber);
! 	Assert(ItemIdIsNormal(lp));
! 
! 	tuple.t_data = (HeapTupleHeader) PageGetItem(page, lp);
! 	tupDesc = RelationGetDescr(seq_rel);
! 	values=palloc(sizeof(Datum)*tupDesc->natts);
! 	isnull=palloc(sizeof(bool)*tupDesc->natts);
! 	heap_deform_tuple(&tuple,tupDesc,values,isnull);
! 	UnlockReleaseBuffer(buf);
! 	relation_close(seq_rel,NoLock);
! 
! 	/**
! 	 * Get a new Relfilenode
! 	 */
! 	RelationSetNewRelfilenode(seq_rel,RecentXmin);
! 
! 	/**
! 	 * Add a tuple to the new refilenode. use the values
! 	 * from the old version.
! 	 */
! 	init_seq_relation(seq_rel,tupDesc,values,isnull,NULL);
! 	pfree(values);
! 	pfree(isnull);
! 	/**
! 	 * reset the sequence
! 	 */
! 	AlterSequenceInternal(seq_relid, options);
! 	relation_close(seq_rel,NoLock);
! 
! }
! static void init_seq_relation(Relation rel,TupleDesc tupDesc,Datum *value,
! 							  bool * null, List * owned_by)
! {
! 	Buffer buf;
! 	Page page;
! 	sequence_magic * sm;
! 	HeapTuple tuple;
! 
! 
! 	tupDesc = RelationGetDescr(rel);
! 		/* Initialize first page of relation with special magic number */
  
  	buf = ReadBuffer(rel, P_NEW);
  	Assert(BufferGetBlockNumber(buf) == 0);
*************** DefineSequence(CreateSeqStmt *seq)
*** 311,317 ****
  	if (owned_by)
  		process_owned_by(rel, owned_by);
  
! 	heap_close(rel, NoLock);
  }
  
  /*
--- 387,393 ----
  	if (owned_by)
  		process_owned_by(rel, owned_by);
  
! 
  }
  
  /*
diff --git a/src/backend/commands/tablecmds.c b/src/backend/commands/tablecmds.c
index c009711..04318f8 100644
*** a/src/backend/commands/tablecmds.c
--- b/src/backend/commands/tablecmds.c
*************** ExecuteTruncate(TruncateStmt *stmt)
*** 816,822 ****
  {
  	List	   *rels = NIL;
  	List	   *relids = NIL;
- 	List	   *seq_relids = NIL;
  	EState	   *estate;
  	ResultRelInfo *resultRelInfos;
  	ResultRelInfo *resultRelInfo;
--- 816,821 ----
*************** ExecuteTruncate(TruncateStmt *stmt)
*** 923,949 ****
  	 */
  	if (stmt->restart_seqs)
  	{
  		foreach(cell, rels)
  		{
  			Relation	rel = (Relation) lfirst(cell);
  			List	   *seqlist = getOwnedSequences(RelationGetRelid(rel));
  			ListCell   *seqcell;
! 
  			foreach(seqcell, seqlist)
  			{
  				Oid			seq_relid = lfirst_oid(seqcell);
! 				Relation	seq_rel;
! 
! 				seq_rel = relation_open(seq_relid, AccessShareLock);
! 
  				/* This check must match AlterSequence! */
  				if (!pg_class_ownercheck(seq_relid, GetUserId()))
  					aclcheck_error(ACLCHECK_NOT_OWNER, ACL_KIND_CLASS,
  								   RelationGetRelationName(seq_rel));
- 
- 				seq_relids = lappend_oid(seq_relids, seq_relid);
- 
  				relation_close(seq_rel, NoLock);
  			}
  		}
  	}
--- 922,944 ----
  	 */
  	if (stmt->restart_seqs)
  	{
+ 
  		foreach(cell, rels)
  		{
  			Relation	rel = (Relation) lfirst(cell);
  			List	   *seqlist = getOwnedSequences(RelationGetRelid(rel));
  			ListCell   *seqcell;
! 			List	   *options = list_make1(makeDefElem("restart", NULL));
  			foreach(seqcell, seqlist)
  			{
  				Oid			seq_relid = lfirst_oid(seqcell);
! 				Relation seq_rel = relation_open(seq_relid, AccessShareLock);
  				/* This check must match AlterSequence! */
  				if (!pg_class_ownercheck(seq_relid, GetUserId()))
  					aclcheck_error(ACLCHECK_NOT_OWNER, ACL_KIND_CLASS,
  								   RelationGetRelationName(seq_rel));
  				relation_close(seq_rel, NoLock);
+ 				ResetSequenceRelation(seq_relid,options);
  			}
  		}
  	}
*************** ExecuteTruncate(TruncateStmt *stmt)
*** 1068,1093 ****
  		heap_close(rel, NoLock);
  	}
  
- 	/*
- 	 * Lastly, restart any owned sequences if we were asked to.  This is done
- 	 * last because it's nontransactional: restarts will not roll back if we
- 	 * abort later.  Hence it's important to postpone them as long as
- 	 * possible.  (This is also a big reason why we locked and
- 	 * permission-checked the sequences beforehand.)
- 	 */
- 	if (stmt->restart_seqs)
- 	{
- 		List	   *options = list_make1(makeDefElem("restart", NULL));
- 
- 		foreach(cell, seq_relids)
- 		{
- 			Oid			seq_relid = lfirst_oid(cell);
- 
- 			AlterSequenceInternal(seq_relid, options);
- 		}
- 	}
  }
  
  /*
   * Check that a given rel is safe to truncate.	Subroutine for ExecuteTruncate
   */
--- 1063,1071 ----
  		heap_close(rel, NoLock);
  	}
  
  }
  
+ 
  /*
   * Check that a given rel is safe to truncate.	Subroutine for ExecuteTruncate
   */
diff --git a/src/include/commands/sequence.h b/src/include/commands/sequence.h
index 5f566f6..3dab6db 100644
*** a/src/include/commands/sequence.h
--- b/src/include/commands/sequence.h
*************** extern Datum lastval(PG_FUNCTION_ARGS);
*** 72,77 ****
--- 72,78 ----
  extern void DefineSequence(CreateSeqStmt *stmt);
  extern void AlterSequence(AlterSeqStmt *stmt);
  extern void AlterSequenceInternal(Oid relid, List *options);
+ extern void ResetSequenceRelation(Oid seq_relid,List * options);
  
  extern void seq_redo(XLogRecPtr lsn, XLogRecord *rptr);
  extern void seq_desc(StringInfo buf, uint8 xl_info, char *rec);
diff --git a/src/test/regress/expected/truncate.out b/src/test/regress/expected/truncate.out
index 7f43df7..169c53a 100644
*** a/src/test/regress/expected/truncate.out
--- b/src/test/regress/expected/truncate.out
*************** SELECT * FROM truncate_a;
*** 398,403 ****
--- 398,415 ----
    2 |  34
  (2 rows)
  
+ BEGIN;
+ TRUNCATE truncate_a RESTART IDENTITY;
+ ROLLBACK;
+ INSERT INTO truncate_a DEFAULT VALUES;
+ SELECT * FROM truncate_a;
+  id | id1 
+ ----+-----
+   1 |  33
+   2 |  34
+   3 |  35
+ (3 rows)
+ 
  DROP TABLE truncate_a;
  SELECT nextval('truncate_a_id1'); -- fail, seq should have been dropped
  ERROR:  relation "truncate_a_id1" does not exist
diff --git a/src/test/regress/sql/truncate.sql b/src/test/regress/sql/truncate.sql
index b348e94..9c2c91a 100644
*** a/src/test/regress/sql/truncate.sql
--- b/src/test/regress/sql/truncate.sql
*************** INSERT INTO truncate_a DEFAULT VALUES;
*** 202,207 ****
--- 202,214 ----
  INSERT INTO truncate_a DEFAULT VALUES;
  SELECT * FROM truncate_a;
  
+ BEGIN;
+ TRUNCATE truncate_a RESTART IDENTITY;
+ ROLLBACK;
+ INSERT INTO truncate_a DEFAULT VALUES;
+ SELECT * FROM truncate_a;
+ 
+ 
  DROP TABLE truncate_a;
  
  SELECT nextval('truncate_a_id1'); -- fail, seq should have been dropped