From 21bb5d22492089a804903a817b8dd5c179643056 Mon Sep 17 00:00:00 2001
From: Justin Pryzby <pryzbyj@telsasoft.com>
Date: Tue, 13 Jul 2021 21:25:48 -0500
Subject: [PATCH v2 1/5] psql: \dn+ to show size of each schema..

See also: 358a897fa, 528ac10c7
---
 src/backend/utils/adt/dbsize.c  | 81 +++++++++++++++++++++++++++++++++
 src/bin/psql/describe.c         |  7 +++
 src/include/catalog/pg_proc.dat | 10 ++++
 3 files changed, 98 insertions(+)

diff --git a/src/backend/utils/adt/dbsize.c b/src/backend/utils/adt/dbsize.c
index d5a7fb13f3..d13e691807 100644
--- a/src/backend/utils/adt/dbsize.c
+++ b/src/backend/utils/adt/dbsize.c
@@ -13,8 +13,10 @@
 
 #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"
@@ -25,6 +27,8 @@
 #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 +836,83 @@ pg_size_bytes(PG_FUNCTION_ARGS)
 	PG_RETURN_INT64(result);
 }
 
+/* Compute the size of a schema (namespace) */
+static int64
+calculate_namespace_size(Oid nspOid)
+{
+	int64		totalsize = 0;
+	AclResult	aclresult;
+	Relation	pg_class;
+	ScanKeyData 	skey;
+	SysScanDesc	scan;
+	HeapTuple	tuple;
+
+	/*
+	 * 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 = pg_namespace_aclcheck(nspOid, GetUserId(), ACL_CREATE);
+		if (aclresult != ACLCHECK_OK)
+			aclcheck_error(aclresult, OBJECT_SCHEMA,
+						   get_namespace_name(nspOid));
+	}
+
+	ScanKeyInit(&skey, Anum_pg_class_relnamespace,
+			BTEqualStrategyNumber, F_OIDEQ, nspOid);
+
+	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 = 1; 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;
+}
+
+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);
+}
+
 /*
  * Get the filenode of a relation
  *
diff --git a/src/bin/psql/describe.c b/src/bin/psql/describe.c
index ba658f731b..76b50b429d 100644
--- a/src/bin/psql/describe.c
+++ b/src/bin/psql/describe.c
@@ -5036,6 +5036,13 @@ 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  (SELECT pg_catalog.pg_size_pretty(pg_namespace_size(n.oid))) AS \"%s\"",
+							  // ",\n  (SELECT pg_catalog.pg_size_pretty(sum(pg_relation_size(oid,fork))) FROM pg_catalog.pg_class c,\n"
+							  // "     (VALUES('main'),('fsm'),('vm'),('init')) AS fork(fork) WHERE c.relnamespace = 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..2d07f66b67 100644
--- a/src/include/catalog/pg_proc.dat
+++ b/src/include/catalog/pg_proc.dat
@@ -7158,6 +7158,16 @@
   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 => '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