From d4a2eb119a1e956fb022ca06c860a114b6477f6b Mon Sep 17 00:00:00 2001
From: Georgios Kokolatos <gkokolatos@pm.me>
Date: Wed, 24 Feb 2021 13:36:44 +0000
Subject: [PATCH v4] Attempt to make dbsize a bit more consistent

---
 src/backend/access/table/tableam.c |   8 +-
 src/backend/utils/adt/dbsize.c     | 230 +++++++++++++++++++----------
 src/bin/psql/describe.c            |  16 +-
 src/include/catalog/pg_proc.dat    |   8 +
 src/test/regress/expected/psql.out |  12 +-
 5 files changed, 183 insertions(+), 91 deletions(-)

diff --git a/src/backend/access/table/tableam.c b/src/backend/access/table/tableam.c
index 5ea5bdd810..af0188177a 100644
--- a/src/backend/access/table/tableam.c
+++ b/src/backend/access/table/tableam.c
@@ -636,10 +636,14 @@ table_block_relation_size(Relation rel, ForkNumber forkNumber)
 	if (forkNumber == InvalidForkNumber)
 	{
 		for (int i = 0; i < MAX_FORKNUM; i++)
-			nblocks += smgrnblocks(rel->rd_smgr, i);
+			nblocks += smgrexists(rel->rd_smgr, i)
+						? smgrnblocks(rel->rd_smgr, i)
+						: 0;
 	}
 	else
-		nblocks = smgrnblocks(rel->rd_smgr, forkNumber);
+		nblocks = smgrexists(rel->rd_smgr, forkNumber)
+					? smgrnblocks(rel->rd_smgr, forkNumber)
+					: 0;
 
 	return nblocks * BLCKSZ;
 }
diff --git a/src/backend/utils/adt/dbsize.c b/src/backend/utils/adt/dbsize.c
index 64cdaa4134..76c7427a41 100644
--- a/src/backend/utils/adt/dbsize.c
+++ b/src/backend/utils/adt/dbsize.c
@@ -15,6 +15,7 @@
 
 #include "access/htup_details.h"
 #include "access/relation.h"
+#include "access/tableam.h"
 #include "catalog/catalog.h"
 #include "catalog/namespace.h"
 #include "catalog/pg_authid.h"
@@ -34,6 +35,8 @@
 /* Divide by two and round towards positive infinity. */
 #define half_rounded(x)   (((x) + ((x) < 0 ? 0 : 1)) / 2)
 
+static int64 calculate_total_relation_size(Relation rel);
+
 /* Return physical size of directory contents, or 0 if dir doesn't exist */
 static int64
 db_dir_size(const char *path)
@@ -307,104 +310,93 @@ calculate_relation_size(RelFileNode *rfn, BackendId backend, ForkNumber forknum)
 	return totalsize;
 }
 
-Datum
-pg_relation_size(PG_FUNCTION_ARGS)
-{
-	Oid			relOid = PG_GETARG_OID(0);
-	text	   *forkName = PG_GETARG_TEXT_PP(1);
-	Relation	rel;
-	int64		size;
-
-	rel = try_relation_open(relOid, AccessShareLock);
-
-	/*
-	 * Before 9.2, we used to throw an error if the relation didn't exist, but
-	 * that makes queries like "SELECT pg_relation_size(oid) FROM pg_class"
-	 * less robust, because while we scan pg_class with an MVCC snapshot,
-	 * someone else might drop the table. It's better to return NULL for
-	 * already-dropped tables than throw an error and abort the whole query.
-	 */
-	if (rel == NULL)
-		PG_RETURN_NULL();
-
-	size = calculate_relation_size(&(rel->rd_node), rel->rd_backend,
-								   forkname_to_number(text_to_cstring(forkName)));
-
-	relation_close(rel, AccessShareLock);
-
-	PG_RETURN_INT64(size);
-}
-
 /*
- * Calculate total on-disk size of a TOAST relation, including its indexes.
- * Must not be applied to non-TOAST relations.
+ * Calculate total on-disk size of a given table,
+ * plus TOAST table if any.
+ * Indexes other than the TOAST table's index are not included.
+ *
+ * Note that this also behaves sanely if applied to a toast table.
  */
 static int64
-calculate_toast_table_size(Oid toastrelid)
+calculate_table_size(Relation rel)
 {
-	int64		size = 0;
-	Relation	toastRel;
+	uint64		size = 0;
 	ForkNumber	forkNum;
-	ListCell   *lc;
-	List	   *indexlist;
-
-	toastRel = relation_open(toastrelid, AccessShareLock);
 
-	/* toast heap size, including FSM and VM size */
+	/*
+	 * table size, including all implemented forks (e.g. FSM, VM for heap AM)
+	 * excluding TOAST relation
+	 */
 	for (forkNum = 0; forkNum <= MAX_FORKNUM; forkNum++)
-		size += calculate_relation_size(&(toastRel->rd_node),
-										toastRel->rd_backend, forkNum);
+		size += table_relation_size(rel, forkNum);
 
-	/* toast index size, including FSM and VM size */
-	indexlist = RelationGetIndexList(toastRel);
-
-	/* Size is calculated using all the indexes available */
-	foreach(lc, indexlist)
+	/*
+	 * Size of toast relation
+	 */
+	if (OidIsValid(rel->rd_rel->reltoastrelid))
 	{
-		Relation	toastIdxRel;
+		Relation	toastRel;
 
-		toastIdxRel = relation_open(lfirst_oid(lc),
-									AccessShareLock);
-		for (forkNum = 0; forkNum <= MAX_FORKNUM; forkNum++)
-			size += calculate_relation_size(&(toastIdxRel->rd_node),
-											toastIdxRel->rd_backend, forkNum);
+		toastRel = relation_open(rel->rd_rel->reltoastrelid, AccessShareLock);
 
-		relation_close(toastIdxRel, AccessShareLock);
+		size += calculate_total_relation_size(toastRel);
+
+		relation_close(toastRel, AccessShareLock);
 	}
-	list_free(indexlist);
-	relation_close(toastRel, AccessShareLock);
 
-	return size;
+	return (int64)size;
 }
 
+
 /*
- * Calculate total on-disk size of a given table,
- * including FSM and VM, plus TOAST table if any.
- * Indexes other than the TOAST table's index are not included.
+ * Calculate on-disk size of a given table for the specified fork
+ * if implemented by the table AM.
  *
- * Note that this also behaves sanely if applied to an index or toast table;
- * those won't have attached toast tables, but they can have multiple forks.
+ * Note that this also behaves sanely if applied to a toast table.
  */
 static int64
-calculate_table_size(Relation rel)
+calculate_table_fork_size(Relation rel, ForkNumber forkNum)
 {
-	int64		size = 0;
-	ForkNumber	forkNum;
+	return (int64)table_relation_size(rel, forkNum);
+}
 
-	/*
-	 * heap size, including FSM and VM
-	 */
-	for (forkNum = 0; forkNum <= MAX_FORKNUM; forkNum++)
-		size += calculate_relation_size(&(rel->rd_node), rel->rd_backend,
-										forkNum);
+Datum
+pg_relation_size(PG_FUNCTION_ARGS)
+{
+	Oid			relOid = PG_GETARG_OID(0);
+	text	   *forkName = PG_GETARG_TEXT_PP(1);
+	Relation	rel;
+	int64		size;
+
+	rel = try_relation_open(relOid, AccessShareLock);
 
 	/*
-	 * Size of toast relation
+	 * Before 9.2, we used to throw an error if the relation didn't exist, but
+	 * that makes queries like "SELECT pg_relation_size(oid) FROM pg_class"
+	 * less robust, because while we scan pg_class with an MVCC snapshot,
+	 * someone else might drop the table. It's better to return NULL for
+	 * already-dropped tables than throw an error and abort the whole query.
 	 */
-	if (OidIsValid(rel->rd_rel->reltoastrelid))
-		size += calculate_toast_table_size(rel->rd_rel->reltoastrelid);
+	if (rel == NULL)
+		PG_RETURN_NULL();
 
-	return size;
+	if (rel->rd_rel->relkind == RELKIND_RELATION ||
+		rel->rd_rel->relkind == RELKIND_TOASTVALUE ||
+		rel->rd_rel->relkind == RELKIND_MATVIEW)
+		size = calculate_table_fork_size(rel,
+										forkname_to_number(text_to_cstring(forkName)));
+	else if (rel->rd_rel->relkind == RELKIND_INDEX)
+		size = calculate_relation_size(&(rel->rd_node), rel->rd_backend,
+									   forkname_to_number(text_to_cstring(forkName)));
+	else
+	{
+		relation_close(rel, AccessShareLock);
+		PG_RETURN_NULL();
+	}
+
+	relation_close(rel, AccessShareLock);
+
+	PG_RETURN_INT64(size);
 }
 
 /*
@@ -415,7 +407,7 @@ calculate_table_size(Relation rel)
 static int64
 calculate_indexes_size(Relation rel)
 {
-	int64		size = 0;
+	uint64		size = 0;
 
 	/*
 	 * Aggregate all indexes on the given relation
@@ -444,7 +436,7 @@ calculate_indexes_size(Relation rel)
 		list_free(index_oids);
 	}
 
-	return size;
+	return (int64)size;
 }
 
 Datum
@@ -452,14 +444,22 @@ pg_table_size(PG_FUNCTION_ARGS)
 {
 	Oid			relOid = PG_GETARG_OID(0);
 	Relation	rel;
-	int64		size;
+	int64		size = 0;
 
 	rel = try_relation_open(relOid, AccessShareLock);
 
 	if (rel == NULL)
 		PG_RETURN_NULL();
 
-	size = calculate_table_size(rel);
+	if (rel->rd_rel->relkind == RELKIND_RELATION ||
+		rel->rd_rel->relkind == RELKIND_TOASTVALUE ||
+		rel->rd_rel->relkind == RELKIND_MATVIEW)
+		size = calculate_table_size(rel);
+	else
+	{
+		relation_close(rel, AccessShareLock);
+		PG_RETURN_NULL();
+	}
 
 	relation_close(rel, AccessShareLock);
 
@@ -485,14 +485,72 @@ pg_indexes_size(PG_FUNCTION_ARGS)
 	PG_RETURN_INT64(size);
 }
 
+Datum
+pg_index_size(PG_FUNCTION_ARGS)
+{
+	Oid			relOid = PG_GETARG_OID(0);
+	Relation	rel;
+	int64		size = 0;
+
+	rel = try_relation_open(relOid, AccessShareLock);
+
+	if (rel == NULL)
+		PG_RETURN_NULL();
+
+	if (rel->rd_rel->relkind == RELKIND_INDEX)
+	{
+		for (ForkNumber forkNum = 0; forkNum <= MAX_FORKNUM; forkNum++)
+			size += calculate_relation_size(&(rel->rd_node),
+											rel->rd_backend,
+											forkNum);
+	}
+	else
+	{
+		relation_close(rel, AccessShareLock);
+		PG_RETURN_NULL();
+	}
+
+	relation_close(rel, AccessShareLock);
+
+	PG_RETURN_INT64(size);
+}
+
+Datum
+pg_sequence_size(PG_FUNCTION_ARGS)
+{
+	Oid			relOid = PG_GETARG_OID(0);
+	Relation	rel;
+	int64		size = 0;
+
+	rel = try_relation_open(relOid, AccessShareLock);
+
+	if (rel == NULL)
+		PG_RETURN_NULL();
+
+	if (rel->rd_rel->relkind != RELKIND_SEQUENCE)
+	{
+		relation_close(rel, AccessShareLock);
+		PG_RETURN_NULL();
+	}
+
+	for (ForkNumber forkNum = 0; forkNum <= MAX_FORKNUM; forkNum++)
+		size += calculate_relation_size(&(rel->rd_node),
+										rel->rd_backend,
+										forkNum);
+
+	relation_close(rel, AccessShareLock);
+
+	PG_RETURN_INT64(size);
+}
+
 /*
  *	Compute the on-disk size of all files for the relation,
- *	including heap data, index data, toast data, FSM, VM.
+ *	including table data, index data, toast data.
  */
 static int64
 calculate_total_relation_size(Relation rel)
 {
-	int64		size;
+	uint64		size;
 
 	/*
 	 * Aggregate the table size, this includes size of the heap, toast and
@@ -505,7 +563,7 @@ calculate_total_relation_size(Relation rel)
 	 */
 	size += calculate_indexes_size(rel);
 
-	return size;
+	return (int64)size;
 }
 
 Datum
@@ -513,14 +571,22 @@ pg_total_relation_size(PG_FUNCTION_ARGS)
 {
 	Oid			relOid = PG_GETARG_OID(0);
 	Relation	rel;
-	int64		size;
+	uint64		size = 0;
 
 	rel = try_relation_open(relOid, AccessShareLock);
 
 	if (rel == NULL)
 		PG_RETURN_NULL();
 
-	size = calculate_total_relation_size(rel);
+	if (rel->rd_rel->relkind == RELKIND_RELATION ||
+		rel->rd_rel->relkind == RELKIND_TOASTVALUE ||
+		rel->rd_rel->relkind == RELKIND_MATVIEW)
+		size = calculate_total_relation_size(rel);
+	else
+	{
+		relation_close(rel, AccessShareLock);
+		PG_RETURN_NULL();
+	}
 
 	relation_close(rel, AccessShareLock);
 
diff --git a/src/bin/psql/describe.c b/src/bin/psql/describe.c
index 20af5a92b4..84f55190a3 100644
--- a/src/bin/psql/describe.c
+++ b/src/bin/psql/describe.c
@@ -3782,10 +3782,24 @@ listTables(const char *tabtypes, const char *pattern, bool verbose, bool showSys
 							  gettext_noop("Access Method"));
 
 		/*
+		 * As of PostgreSQL 14, do not use pg_table_size() for indexes and
+		 * sequences as it does not behave sanely for those.
+		 *
 		 * As of PostgreSQL 9.0, use pg_table_size() to show a more accurate
 		 * size of a table, including FSM, VM and TOAST tables.
 		 */
-		if (pset.sversion >= 90000)
+		if (pset.sversion >= 140000)
+			appendPQExpBuffer(&buf,
+							  ",\n CASE"
+							  " WHEN c.relkind in ("CppAsString2(RELKIND_INDEX)") THEN"
+							  "		pg_catalog.pg_size_pretty(pg_catalog.pg_index_size(c.oid))"
+							  " WHEN c.relkind in ("CppAsString2(RELKIND_SEQUENCE)") THEN"
+							  "		pg_catalog.pg_size_pretty(pg_catalog.pg_sequence_size(c.oid))"
+							  " ELSE"
+							  " 	pg_catalog.pg_size_pretty(pg_catalog.pg_table_size(c.oid))"
+							  " END as \"%s\"",
+							  gettext_noop("Size"));
+		else if (pset.sversion >= 90000)
 			appendPQExpBuffer(&buf,
 							  ",\n  pg_catalog.pg_size_pretty(pg_catalog.pg_table_size(c.oid)) as \"%s\"",
 							  gettext_noop("Size"));
diff --git a/src/include/catalog/pg_proc.dat b/src/include/catalog/pg_proc.dat
index 1487710d59..0a953ac3e9 100644
--- a/src/include/catalog/pg_proc.dat
+++ b/src/include/catalog/pg_proc.dat
@@ -7142,6 +7142,14 @@
   descr => 'disk space usage for all indexes attached to the specified table',
   proname => 'pg_indexes_size', provolatile => 'v', prorettype => 'int8',
   proargtypes => 'regclass', prosrc => 'pg_indexes_size' },
+{ oid => '4543',
+  descr => 'disk space usage for the specified index',
+  proname => 'pg_index_size', provolatile => 'v', prorettype => 'int8',
+  proargtypes => 'regclass', prosrc => 'pg_index_size' },
+{ oid => '4544',
+  descr => 'disk space usage for the specified sequence',
+  proname => 'pg_sequence_size', provolatile => 'v', prorettype => 'int8',
+  proargtypes => 'regclass', prosrc => 'pg_sequence_size' },
 { oid => '2999', descr => 'filenode identifier of relation',
   proname => 'pg_relation_filenode', provolatile => 's', prorettype => 'oid',
   proargtypes => 'regclass', prosrc => 'pg_relation_filenode' },
diff --git a/src/test/regress/expected/psql.out b/src/test/regress/expected/psql.out
index 7204fdb0b4..70d7a594dc 100644
--- a/src/test/regress/expected/psql.out
+++ b/src/test/regress/expected/psql.out
@@ -2851,7 +2851,7 @@ Access method: heap
  tableam_display | mat_view_heap_psql | materialized view | regress_display_role | permanent   | heap_psql     | 0 bytes | 
  tableam_display | tbl_heap           | table             | regress_display_role | permanent   | heap          | 0 bytes | 
  tableam_display | tbl_heap_psql      | table             | regress_display_role | permanent   | heap_psql     | 0 bytes | 
- tableam_display | view_heap_psql     | view              | regress_display_role | permanent   |               | 0 bytes | 
+ tableam_display | view_heap_psql     | view              | regress_display_role | permanent   |               |         | 
 (4 rows)
 
 \dt+
@@ -2871,10 +2871,10 @@ Access method: heap
 
 -- But not for views and sequences.
 \dv+
-                                          List of relations
-     Schema      |      Name      | Type |        Owner         | Persistence |  Size   | Description 
------------------+----------------+------+----------------------+-------------+---------+-------------
- tableam_display | view_heap_psql | view | regress_display_role | permanent   | 0 bytes | 
+                                         List of relations
+     Schema      |      Name      | Type |        Owner         | Persistence | Size | Description 
+-----------------+----------------+------+----------------------+-------------+------+-------------
+ tableam_display | view_heap_psql | view | regress_display_role | permanent   |      | 
 (1 row)
 
 \set HIDE_TABLEAM on
@@ -2885,7 +2885,7 @@ Access method: heap
  tableam_display | mat_view_heap_psql | materialized view | regress_display_role | permanent   | 0 bytes | 
  tableam_display | tbl_heap           | table             | regress_display_role | permanent   | 0 bytes | 
  tableam_display | tbl_heap_psql      | table             | regress_display_role | permanent   | 0 bytes | 
- tableam_display | view_heap_psql     | view              | regress_display_role | permanent   | 0 bytes | 
+ tableam_display | view_heap_psql     | view              | regress_display_role | permanent   |         | 
 (4 rows)
 
 RESET ROLE;
-- 
2.25.1

