From d7f5446c0aa1a12a512482d2ed121eb44f6e2f59 Mon Sep 17 00:00:00 2001
From: Justin Pryzby <pryzbyj@telsasoft.com>
Date: Tue, 13 Jul 2021 21:25:48 -0500
Subject: [PATCH v3 1/4] psql: show size of schemas and AMs in \dn+ and \dA+..

See also: 358a897fa, 528ac10c7

Add pg_am_size() and psql \dA++
---
 src/backend/utils/adt/dbsize.c  | 138 ++++++++++++++++++++++++++++++++
 src/bin/psql/describe.c         |  12 ++-
 src/include/catalog/pg_proc.dat |  19 +++++
 3 files changed, 168 insertions(+), 1 deletion(-)

diff --git a/src/backend/utils/adt/dbsize.c b/src/backend/utils/adt/dbsize.c
index d5a7fb13f3..62273d4f15 100644
--- a/src/backend/utils/adt/dbsize.c
+++ b/src/backend/utils/adt/dbsize.c
@@ -13,18 +13,23 @@
 
 #include <sys/stat.h>
 
+#include "access/genam.h"
 #include "access/htup_details.h"
 #include "access/relation.h"
+#include "access/table.h"
 #include "catalog/catalog.h"
 #include "catalog/namespace.h"
 #include "catalog/pg_authid.h"
 #include "catalog/pg_tablespace.h"
 #include "commands/dbcommands.h"
+#include "commands/defrem.h"
 #include "commands/tablespace.h"
 #include "miscadmin.h"
 #include "storage/fd.h"
 #include "utils/acl.h"
 #include "utils/builtins.h"
+#include "utils/fmgroids.h"
+#include "utils/lsyscache.h"
 #include "utils/numeric.h"
 #include "utils/rel.h"
 #include "utils/relfilenodemap.h"
@@ -832,6 +837,139 @@ pg_size_bytes(PG_FUNCTION_ARGS)
 	PG_RETURN_INT64(result);
 }
 
+/* Return the sum of size of relations for which the given attribute has the specified OID */
+static int64
+calculate_size_attvalue(int attnum, Oid attval)
+{
+	int64		totalsize = 0;
+	ScanKeyData 	skey;
+	Relation	pg_class;
+	SysScanDesc	scan;
+	HeapTuple	tuple;
+
+	ScanKeyInit(&skey, attnum,
+			BTEqualStrategyNumber, F_OIDEQ, attval);
+
+	pg_class = table_open(RelationRelationId, AccessShareLock);
+	scan = systable_beginscan(pg_class, InvalidOid, false, NULL, 1, &skey);
+	while (HeapTupleIsValid(tuple = systable_getnext(scan)))
+	{
+		Relation rel;
+		Form_pg_class classtuple = (Form_pg_class) GETSTRUCT(tuple);
+
+		rel = try_relation_open(classtuple->oid, AccessShareLock);
+		if (!rel)
+			continue;
+
+		for (int forkNum = 0; forkNum <= MAX_FORKNUM; forkNum++)
+			totalsize += calculate_relation_size(&rel->rd_node, rel->rd_backend, forkNum);
+
+		relation_close(rel, AccessShareLock);
+	}
+
+	systable_endscan(scan);
+	table_close(pg_class, AccessShareLock);
+	return totalsize;
+}
+
+/* Compute the size of relations in a schema (namespace) */
+static int64
+calculate_namespace_size(Oid nspOid)
+{
+	/*
+	 * User must be a member of pg_read_all_stats or have CREATE privilege for
+	 * target namespace.
+	 */
+	if (!is_member_of_role(GetUserId(), ROLE_PG_READ_ALL_STATS))
+	{
+		AclResult aclresult;
+		aclresult = pg_namespace_aclcheck(nspOid, GetUserId(), ACL_CREATE);
+		if (aclresult != ACLCHECK_OK)
+			aclcheck_error(aclresult, OBJECT_SCHEMA,
+						   get_namespace_name(nspOid));
+	}
+
+	return calculate_size_attvalue(Anum_pg_class_relnamespace, nspOid);
+}
+
+Datum
+pg_namespace_size_oid(PG_FUNCTION_ARGS)
+{
+	int64		size;
+	Oid			nspOid = PG_GETARG_OID(0);
+
+	size = calculate_namespace_size(nspOid);
+
+	if (size < 0)
+		PG_RETURN_NULL();
+
+	PG_RETURN_INT64(size);
+}
+
+Datum
+pg_namespace_size_name(PG_FUNCTION_ARGS)
+{
+	int64		size;
+	Name		nspName = PG_GETARG_NAME(0);
+	Oid			nspOid = get_namespace_oid(NameStr(*nspName), false);
+
+	size = calculate_namespace_size(nspOid);
+
+	if (size < 0)
+		PG_RETURN_NULL();
+
+	PG_RETURN_INT64(size);
+}
+
+/* Compute the size of relations using an access method */
+static int64
+calculate_am_size(Oid amOid)
+{
+	/*
+	 * User must be a member of pg_read_all_stats or have CREATE privilege for
+	 * target namespace.
+	 */
+	if (!is_member_of_role(GetUserId(), ROLE_PG_READ_ALL_STATS))
+	{
+		AclResult aclresult;
+		aclresult = pg_namespace_aclcheck(amOid, GetUserId(), ACL_CREATE);
+		if (aclresult != ACLCHECK_OK)
+			aclcheck_error(aclresult, OBJECT_SCHEMA,
+						   get_am_name(amOid));
+	}
+
+	return calculate_size_attvalue(Anum_pg_class_relam, amOid);
+}
+
+Datum
+pg_am_size_oid(PG_FUNCTION_ARGS)
+{
+	int64		size;
+	Oid			amOid = PG_GETARG_OID(0);
+
+	size = calculate_am_size(amOid);
+
+	if (size < 0)
+		PG_RETURN_NULL();
+
+	PG_RETURN_INT64(size);
+}
+
+Datum
+pg_am_size_name(PG_FUNCTION_ARGS)
+{
+	int64		size;
+	Name		amName = PG_GETARG_NAME(0);
+	Oid			amOid = get_am_oid(NameStr(*amName), false);
+
+	size = calculate_am_size(amOid);
+
+	if (size < 0)
+		PG_RETURN_NULL();
+
+	PG_RETURN_INT64(size);
+}
+
 /*
  * Get the filenode of a relation
  *
diff --git a/src/bin/psql/describe.c b/src/bin/psql/describe.c
index ba658f731b..fe5d88a2b7 100644
--- a/src/bin/psql/describe.c
+++ b/src/bin/psql/describe.c
@@ -152,7 +152,7 @@ describeAccessMethods(const char *pattern, bool verbose)
 	PQExpBufferData buf;
 	PGresult   *res;
 	printQueryOpt myopt = pset.popt;
-	static const bool translate_columns[] = {false, true, false, false};
+	static const bool translate_columns[] = {false, true, false, false, false};
 
 	if (pset.sversion < 90600)
 	{
@@ -184,6 +184,11 @@ describeAccessMethods(const char *pattern, bool verbose)
 						  "  pg_catalog.obj_description(oid, 'pg_am') AS \"%s\"",
 						  gettext_noop("Handler"),
 						  gettext_noop("Description"));
+
+		if (pset.sversion >= 150000)
+			appendPQExpBuffer(&buf,
+							  ",\n  pg_catalog.pg_size_pretty(pg_catalog.pg_am_size(oid)) AS \"%s\"",
+							  gettext_noop("Size"));
 	}
 
 	appendPQExpBufferStr(&buf,
@@ -5036,6 +5041,11 @@ listSchemas(const char *pattern, bool verbose, bool showSystem)
 		appendPQExpBuffer(&buf,
 						  ",\n  pg_catalog.obj_description(n.oid, 'pg_namespace') AS \"%s\"",
 						  gettext_noop("Description"));
+
+		if (verbose && pset.sversion >= 150000)
+			appendPQExpBuffer(&buf,
+							  ",\n  pg_catalog.pg_size_pretty(pg_namespace_size(n.oid)) AS \"%s\"",
+							  gettext_noop("Size"));
 	}
 
 	appendPQExpBufferStr(&buf,
diff --git a/src/include/catalog/pg_proc.dat b/src/include/catalog/pg_proc.dat
index fde251fa4f..4518aad200 100644
--- a/src/include/catalog/pg_proc.dat
+++ b/src/include/catalog/pg_proc.dat
@@ -7158,6 +7158,25 @@
   descr => 'total disk space usage for the specified tablespace',
   proname => 'pg_tablespace_size', provolatile => 'v', prorettype => 'int8',
   proargtypes => 'name', prosrc => 'pg_tablespace_size_name' },
+
+{ oid => '9410',
+  descr => 'total disk space usage for the specified namespace',
+  proname => 'pg_namespace_size', provolatile => 'v', prorettype => 'int8',
+  proargtypes => 'oid', prosrc => 'pg_namespace_size_oid' },
+{ oid => '9411',
+  descr => 'total disk space usage for the specified namespace',
+  proname => 'pg_namespace_size', provolatile => 'v', prorettype => 'int8',
+  proargtypes => 'name', prosrc => 'pg_namespace_size_name' },
+
+{ oid => '9412',
+  descr => 'total disk space usage for the specified access method',
+  proname => 'pg_am_size', provolatile => 'v', prorettype => 'int8',
+  proargtypes => 'oid', prosrc => 'pg_am_size_oid' },
+{ oid => '9413',
+  descr => 'total disk space usage for the specified access method',
+  proname => 'pg_am_size', provolatile => 'v', prorettype => 'int8',
+  proargtypes => 'name', prosrc => 'pg_am_size_name' },
+
 { oid => '2324', descr => 'total disk space usage for the specified database',
   proname => 'pg_database_size', provolatile => 'v', prorettype => 'int8',
   proargtypes => 'oid', prosrc => 'pg_database_size_oid' },
-- 
2.17.0

