Add pg_ensure_direct_toast for in-place legacy TOAST table upgrade

When a database cluster is upgraded from a pre-Direct-TOAST PostgreSQL
version or when an existing table with a legacy 3-column TOAST table is
switched to Direct TOAST, the TOAST table may lack the 'chunk_tids' and
'chunk_tid_offsets' columns and the partial index predicate.

This commit introduces:
1. ensure_direct_toast(Oid relid): Performs in-place catalog metadata
   upgrades to add the missing columns in pg_attribute, updates
   pg_class.relnatts = 5, and marks the TOAST index as partial with
   WHERE chunk_id IS NOT NULL in pg_index.
2. pg_ensure_direct_toast(regclass): SQL function callable by DBAs.
3. Automatic upgrade hook in ATExecSetRelOptions() when setting
   toast_flavour = 'direct'.
4. Write-time guard in toast_save_datum_direct() rejecting writes to
   legacy 3-column TOAST tables with an informative hint.
5. Regression test suite and documentation updates.

Author: postsql

doc/src/sgml/storage.sgml

@@ -484,6 +484,86 @@ Furthermore, <literal>plain</literal> and <literal>direct</literal> <acronym>TOA
 pointers can seamlessly coexist within the same table, allowing tables upgraded from
 older <productname>PostgreSQL</productname> versions to be read transparently and
 gradually adopt Direct <acronym>TOAST</acronym> for new writes.
+When upgrading a pre-existing table that was created prior to Direct <acronym>TOAST</acronym>
+support (having a legacy 3-column <acronym>TOAST</acronym> table), executing
+<command>ALTER TABLE ... SET (toast_flavour = 'direct')</command> automatically
+upgrades the <acronym>TOAST</acronym> table in-place to the 5-column format with a
+partial index.  Administrators can also explicitly upgrade any table or <acronym>TOAST</acronym>
+table by calling <function>pg_ensure_direct_toast(<type>regclass</type>)</function>.
+</para>
+
+<para>
+To upgrade all user tables' <acronym>TOAST</acronym> tables in a quiet database in bulk,
+the following query can be used:
+<programlisting>
+SELECT c.oid::regclass AS tablename,
+       pg_ensure_direct_toast(c.reltoastrelid::regclass)
+  FROM pg_class AS c
+ WHERE c.reltoastrelid != 0
+   AND EXISTS (SELECT FROM pg_stat_user_tables AS t WHERE t.relid = c.oid);
+</programlisting>
+</para>
+
+<para>
+For a busy production database, tables can be upgraded one by one with lock timeouts
+and growing wait times between retries to avoid lock contention:
+<programlisting>
+DO $$
+DECLARE
+  lock_timeout_ms int := 100;
+  max_attempts int := 10;
+  tablename text;
+  toastoid regclass;
+  attempt_nr int;
+  update_completed boolean;
+  attempts text;
+  any_found boolean := false;
+BEGIN
+  PERFORM set_config('lock_timeout', lock_timeout_ms || 'ms', false);
+
+  FOR tablename, toastoid IN
+    SELECT c.oid::regclass::text,
+           c.reltoastrelid::regclass
+      FROM pg_class AS c
+     WHERE EXISTS (SELECT FROM pg_stat_user_tables AS t WHERE t.relid = c.oid)
+       AND (SELECT count(*)
+              FROM pg_attribute AS a
+             WHERE a.attrelid = c.reltoastrelid) = 3
+  LOOP
+    any_found := true;
+    update_completed := false;
+    FOR attempt_nr IN 1..max_attempts LOOP
+      BEGIN
+        PERFORM pg_ensure_direct_toast(toastoid);
+        update_completed := true;
+        EXIT;
+      EXCEPTION
+        WHEN lock_not_available THEN
+          PERFORM pg_sleep(0.1 * attempt_nr); /* sleep a little longer each time */
+        WHEN OTHERS THEN
+          RAISE WARNING 'Error updating % to Direct TOAST: %', tablename, SQLERRM;
+          EXIT;
+      END;
+    END LOOP;
+
+    IF update_completed THEN
+      IF attempt_nr > 1 THEN
+        attempts := format(' after %s attempts', attempt_nr);
+      ELSE
+        attempts := '';
+      END IF;
+      RAISE INFO 'Table % updated to Direct TOAST%', tablename, attempts;
+    ELSE
+      RAISE WARNING 'Timeout waiting to update table % to Direct TOAST', tablename;
+    END IF;
+  END LOOP;
+
+  IF NOT any_found THEN
+    RAISE INFO 'No tables needed update to use Direct TOAST';
+  END IF;
+END;
+$$;
+</programlisting>
 </para>
 
 <para>

src/backend/access/common/toast_internals.c

@@ -996,6 +996,15 @@ toast_save_datum_direct(Relation rel, Datum value,
 	state.chunk_seq = 0;
 	state.max_chunk_size = TOAST_MAX_CHUNK_SIZE(TupleDescAttr(state.toasttupDesc, 0)->atttypid);
 
+	if (state.toasttupDesc->natts < 5)
+		ereport(ERROR,
+				(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+				 errmsg("cannot write direct TOAST datum to legacy TOAST table \"%s\"",
+						RelationGetRelationName(state.toastrel)),
+				 errhint("Run \"SELECT pg_ensure_direct_toast('%s'::regclass);\" or \"ALTER TABLE %s SET (toast_flavour = 'direct');\" to upgrade the TOAST table.",
+						 RelationGetRelationName(rel),
+						 RelationGetRelationName(rel))));
+
 	if (VARATT_IS_SHORT(dval))
 	{
 		data_p = VARDATA_SHORT(dval);

src/backend/catalog/toasting.c

@@ -27,10 +27,15 @@
 #include "catalog/pg_am.h"
 #include "catalog/pg_namespace.h"
 #include "catalog/pg_opclass.h"
+#include "catalog/pg_type.h"
 #include "catalog/toasting.h"
 #include "miscadmin.h"
 #include "nodes/makefuncs.h"
+#include "nodes/nodeFuncs.h"
+#include "utils/builtins.h"
 #include "utils/fmgroids.h"
+#include "utils/inval.h"
+#include "utils/lsyscache.h"
 #include "utils/rel.h"
 #include "utils/lsyscache.h"
 #include "utils/syscache.h"
@@ -514,3 +519,176 @@ needs_toast_table(Relation rel)
 	/* Otherwise, let the AM decide. */
 	return table_relation_needs_toast_table(rel);
 }
+
+/*
+ * ensure_direct_toast
+ *
+ * Upgrades a legacy (3-column, full index) TOAST table in-place to support
+ * the Direct TOAST format.
+ * Accepts either the parent table's OID or the TOAST table's OID.
+ */
+void
+ensure_direct_toast(Oid relid)
+{
+	Relation	targetrel;
+	Relation	toastrel;
+	Oid			toastrelid;
+	Relation	pg_class_rel;
+	Relation	pg_attribute_rel;
+	Relation	pg_index_rel;
+	HeapTuple	tuple;
+	HeapTuple	newtuple;
+	TupleDesc	td;
+	Oid			toastIndexOid = InvalidOid;
+	List	   *indexoids;
+
+	targetrel = table_open(relid, AccessShareLock);
+
+	if (targetrel->rd_rel->relkind == RELKIND_TOASTVALUE)
+	{
+		toastrelid = relid;
+		table_close(targetrel, AccessShareLock);
+	}
+	else if (targetrel->rd_rel->relkind == RELKIND_RELATION ||
+			 targetrel->rd_rel->relkind == RELKIND_MATVIEW)
+	{
+		toastrelid = targetrel->rd_rel->reltoastrelid;
+		table_close(targetrel, AccessShareLock);
+
+		if (!OidIsValid(toastrelid))
+			ereport(ERROR,
+					(errcode(ERRCODE_WRONG_OBJECT_TYPE),
+					 errmsg("table \"%s\" does not have a TOAST table",
+							get_rel_name(relid))));
+	}
+	else
+	{
+		table_close(targetrel, AccessShareLock);
+		ereport(ERROR,
+				(errcode(ERRCODE_WRONG_OBJECT_TYPE),
+				 errmsg("relation \"%s\" is not a table or TOAST table",
+						get_rel_name(relid))));
+	}
+
+	/* Open the toast relation with exclusive lock */
+	toastrel = table_open(toastrelid, AccessExclusiveLock);
+
+	if (toastrel->rd_rel->relkind != RELKIND_TOASTVALUE)
+		elog(ERROR, "relation %u is not a TOAST table", toastrelid);
+
+	/*
+	 * Step 1: Add missing columns chunk_tids and chunk_tid_offsets if needed.
+	 */
+	if (toastrel->rd_att->natts < 5)
+	{
+		td = CreateTemplateTupleDesc(2);
+		TupleDescInitEntry(td, (AttrNumber) 1,
+						   "chunk_tids",
+						   TIDARRAYOID,
+						   -1, 0);
+		TupleDescInitEntry(td, (AttrNumber) 2,
+						   "chunk_tid_offsets",
+						   INT8ARRAYOID,
+						   -1, 0);
+		TupleDescAttr(td, 0)->attnum = 4;
+		TupleDescAttr(td, 0)->attstorage = TYPSTORAGE_PLAIN;
+		TupleDescAttr(td, 0)->attcompression = InvalidCompressionMethod;
+		TupleDescAttr(td, 1)->attnum = 5;
+		TupleDescAttr(td, 1)->attstorage = TYPSTORAGE_PLAIN;
+		TupleDescAttr(td, 1)->attcompression = InvalidCompressionMethod;
+
+		populate_compact_attribute(td, 0);
+		populate_compact_attribute(td, 1);
+		TupleDescFinalize(td);
+
+		pg_attribute_rel = table_open(AttributeRelationId, RowExclusiveLock);
+		InsertPgAttributeTuples(pg_attribute_rel, td, toastrelid, NULL, NULL);
+		table_close(pg_attribute_rel, RowExclusiveLock);
+		FreeTupleDesc(td);
+
+		/* Update pg_class.relnatts = 5 */
+		pg_class_rel = table_open(RelationRelationId, RowExclusiveLock);
+		tuple = SearchSysCacheCopy1(RELOID, ObjectIdGetDatum(toastrelid));
+		if (!HeapTupleIsValid(tuple))
+			elog(ERROR, "cache lookup failed for relation %u", toastrelid);
+		((Form_pg_class) GETSTRUCT(tuple))->relnatts = 5;
+		CatalogTupleUpdate(pg_class_rel, &tuple->t_self, tuple);
+		heap_freetuple(tuple);
+		table_close(pg_class_rel, RowExclusiveLock);
+	}
+
+	/*
+	 * Step 2: Ensure the TOAST unique index is partial (WHERE chunk_id IS NOT NULL).
+	 */
+	indexoids = RelationGetIndexList(toastrel);
+	if (indexoids != NIL)
+		toastIndexOid = linitial_oid(indexoids);
+	list_free(indexoids);
+
+	if (OidIsValid(toastIndexOid))
+	{
+		pg_index_rel = table_open(IndexRelationId, RowExclusiveLock);
+		tuple = SearchSysCacheCopy1(INDEXRELID, ObjectIdGetDatum(toastIndexOid));
+		if (HeapTupleIsValid(tuple))
+		{
+			bool		isnull;
+
+			(void) SysCacheGetAttr(INDEXRELID, tuple, Anum_pg_index_indpred, &isnull);
+
+			if (isnull)
+			{
+				Datum		values[Natts_pg_index];
+				bool		nulls[Natts_pg_index];
+				bool		replaces[Natts_pg_index];
+				NullTest   *ntest;
+				char	   *pred_str;
+
+				memset(values, 0, sizeof(values));
+				memset(nulls, 0, sizeof(nulls));
+				memset(replaces, 0, sizeof(replaces));
+
+				ntest = makeNode(NullTest);
+				ntest->arg = (Expr *) makeVar(1, 1, OIDOID, -1, InvalidOid, 0);
+				ntest->nulltesttype = IS_NOT_NULL;
+				ntest->argisrow = false;
+				ntest->location = -1;
+
+				pred_str = nodeToString(list_make1(ntest));
+
+				values[Anum_pg_index_indisprimary - 1] = BoolGetDatum(false);
+				replaces[Anum_pg_index_indisprimary - 1] = true;
+
+				values[Anum_pg_index_indpred - 1] = CStringGetTextDatum(pred_str);
+				replaces[Anum_pg_index_indpred - 1] = true;
+				nulls[Anum_pg_index_indpred - 1] = false;
+
+				newtuple = heap_modify_tuple(tuple, RelationGetDescr(pg_index_rel),
+											 values, nulls, replaces);
+				CatalogTupleUpdate(pg_index_rel, &tuple->t_self, newtuple);
+				heap_freetuple(newtuple);
+				pfree(pred_str);
+
+				CacheInvalidateRelcacheByRelid(toastIndexOid);
+			}
+			heap_freetuple(tuple);
+		}
+		table_close(pg_index_rel, RowExclusiveLock);
+	}
+
+	CacheInvalidateRelcacheByRelid(toastrelid);
+	CommandCounterIncrement();
+
+	table_close(toastrel, AccessExclusiveLock);
+}
+
+/*
+ * SQL-callable function
+ */
+Datum
+pg_ensure_direct_toast(PG_FUNCTION_ARGS)
+{
+	Oid			relid = PG_GETARG_OID(0);
+
+	ensure_direct_toast(relid);
+	PG_RETURN_VOID();
+}

src/backend/commands/tablecmds.c

@@ -17443,6 +17443,19 @@ ATExecSetRelOptions(Relation rel, List *defList, AlterTableType operation,
 
 	ReleaseSysCache(tuple);
 
+	/* If toast_flavour is being set to 'direct', ensure the TOAST table supports it */
+	if (OidIsValid(rel->rd_rel->reltoastrelid) && newOptions != (Datum) 0 &&
+		(rel->rd_rel->relkind == RELKIND_RELATION || rel->rd_rel->relkind == RELKIND_MATVIEW))
+	{
+		StdRdOptions *opts = (StdRdOptions *) heap_reloptions(rel->rd_rel->relkind, newOptions, false);
+
+		if (opts && opts->toast_flavour == TOAST_FLAVOUR_DIRECT)
+			ensure_direct_toast(rel->rd_rel->reltoastrelid);
+
+		if (opts)
+			pfree(opts);
+	}
+
 	/* repeat the whole exercise for the toast table, if there's one */
 	if (OidIsValid(rel->rd_rel->reltoastrelid))
 	{

src/include/catalog/pg_proc.dat

@@ -12778,4 +12778,8 @@
   proname => 'hashoid8extended', prorettype => 'int8',
   proargtypes => 'oid8 int8', prosrc => 'hashoid8extended' },
 
+{ oid => '9225', descr => 'ensure a TOAST table supports Direct TOAST format',
+  proname => 'pg_ensure_direct_toast', provolatile => 'v', prorettype => 'void',
+  proargtypes => 'regclass', prosrc => 'pg_ensure_direct_toast' },
+
 ]

src/include/catalog/toasting.h

@@ -26,5 +26,6 @@ extern void AlterTableCreateToastTable(Oid relOid, Datum reloptions,
 									   LOCKMODE lockmode);
 extern void BootstrapToastTable(char *relName,
 								Oid toastOid, Oid toastIndexOid);
+extern void ensure_direct_toast(Oid relid);
 
 #endif							/* TOASTING_H */