From 501b1fb7df70503d392101737dc8a0e8a8788898 Mon Sep 17 00:00:00 2001
From: Brad Jorsch <bjorsch@wikimedia.org>
Date: Sun, 18 Mar 2018 13:23:58 -0400
Subject: [PATCH] tests: Reset Postgres sequences when cloning and truncating

This improves the repeatability of the unit tests by making the ID
values generated depend less on what previous tests might have done.

It also prevents tests from using up sequence numbers for the live DB's
tables.

Change-Id: Iaa8ae1e5cef4b9099bd1b4b8fc806f5af372a7ff
---
 .../libs/rdbms/database/DatabasePostgres.php  | 72 +++++++++++++++++--
 tests/phpunit/MediaWikiTestCase.php           |  5 ++
 2 files changed, 73 insertions(+), 4 deletions(-)

diff --git a/includes/libs/rdbms/database/DatabasePostgres.php b/includes/libs/rdbms/database/DatabasePostgres.php
index 32ea37592f..525d308e9a 100644
--- a/includes/libs/rdbms/database/DatabasePostgres.php
+++ b/includes/libs/rdbms/database/DatabasePostgres.php
@@ -786,11 +786,75 @@ __INDEXATTR__;
 	public function duplicateTableStructure(
 		$oldName, $newName, $temporary = false, $fname = __METHOD__
 	) {
-		$newName = $this->addIdentifierQuotes( $newName );
-		$oldName = $this->addIdentifierQuotes( $oldName );
+		$newNameE = $this->addIdentifierQuotes( $newName );
+		$oldNameE = $this->addIdentifierQuotes( $oldName );
+
+		$ret = $this->query( 'CREATE ' . ( $temporary ? 'TEMPORARY ' : '' ) . " TABLE $newNameE " .
+			"(LIKE $oldNameE INCLUDING DEFAULTS INCLUDING INDEXES)", $fname );
+		if ( !$ret ) {
+			return $ret;
+		}
+
+		$res = $this->query( 'SELECT attname FROM pg_class c'
+			. ' JOIN pg_namespace n ON (n.oid = c.relnamespace)'
+			. ' JOIN pg_attribute a ON (a.attrelid = c.oid)'
+			. ' JOIN pg_attrdef d ON (c.oid=d.adrelid and a.attnum=d.adnum)'
+			. ' WHERE relkind = \'r\''
+			. ' AND nspname = ' . $this->addQuotes( $this->getCoreSchema() )
+			. ' AND relname = ' . $this->addQuotes( $oldName )
+			. ' AND adsrc LIKE \'nextval(%\'',
+			$fname
+		);
+		$row = $this->fetchObject( $res );
+		if ( $row ) {
+			$field = $row->attname;
+			$newSeq = "{$newName}_{$field}_seq";
+			$fieldE = $this->addIdentifierQuotes( $field );
+			$newSeqE = $this->addIdentifierQuotes( $newSeq );
+			$newSeqQ = $this->addQuotes( $newSeq );
+			$this->query( 'CREATE ' . ( $temporary ? 'TEMPORARY ' : '' ) . " SEQUENCE $newSeqE", $fname );
+			$this->query(
+				"ALTER TABLE $newNameE ALTER COLUMN $fieldE SET DEFAULT nextval({$newSeqQ}::regclass)",
+				$fname
+			);
+		}
 
-		return $this->query( 'CREATE ' . ( $temporary ? 'TEMPORARY ' : '' ) . " TABLE $newName " .
-			"(LIKE $oldName INCLUDING DEFAULTS INCLUDING INDEXES)", $fname );
+		return $ret;
+	}
+
+	public function resetSequenceForTable( $table, $fname = __METHOD__ ) {
+		$table = $this->tableName( $table, 'raw' );
+		foreach ( $this->getCoreSchemas() as $schema ) {
+			$res = $this->query(
+				'SELECT c.oid FROM pg_class c JOIN pg_namespace n ON (n.oid = c.relnamespace)'
+				. ' WHERE relkind = \'r\''
+				. ' AND nspname = ' . $this->addQuotes( $schema )
+				. ' AND relname = ' . $this->addQuotes( $table ),
+				$fname
+			);
+			if ( !$res || !$this->numRows( $res ) ) {
+				continue;
+			}
+
+			$oid = $this->fetchObject( $res )->oid;
+			$res = $this->query( 'SELECT adsrc FROM pg_attribute a'
+				. ' JOIN pg_attrdef d ON (a.attrelid=d.adrelid and a.attnum=d.adnum)'
+				. " WHERE a.attrelid = $oid"
+				. ' AND adsrc LIKE \'nextval(%\'',
+				$fname
+			);
+			$row = $this->fetchObject( $res );
+			if ( $row ) {
+				$this->query(
+					'SELECT ' . preg_replace( '/^nextval\((.+)\)$/', 'setval($1,1,false)', $row->adsrc ),
+					$fname
+				);
+				return true;
+			}
+			return false;
+		}
+
+		return false;
 	}
 
 	public function listTables( $prefix = null, $fname = __METHOD__ ) {
diff --git a/tests/phpunit/MediaWikiTestCase.php b/tests/phpunit/MediaWikiTestCase.php
index 0d2b788b22..ffba861376 100644
--- a/tests/phpunit/MediaWikiTestCase.php
+++ b/tests/phpunit/MediaWikiTestCase.php
@@ -1539,6 +1539,11 @@ abstract class MediaWikiTestCase extends PHPUnit\Framework\TestCase {
 					$db->delete( $tbl, '*', __METHOD__ );
 				}
 
+				if ( $db->getType() === 'postgres' ) {
+					// Reset the table's sequence too.
+					$db->resetSequenceForTable( $tbl, __METHOD__ );
+				}
+
 				if ( $tbl === 'page' ) {
 					// Forget about the pages since they don't
 					// exist in the DB.
-- 
2.20.1