From 528ac10c7c73883f5b12db2f7bd627652f869bab Mon Sep 17 00:00:00 2001
From: Bruce Momjian <bruce@momjian.us>
Date: Thu, 2 Sep 2004 00:55:22 +0000
Subject: [PATCH] The current implementation of dbsize doesn't handle tables in
 tablespaces correctly, and is quite restricted on objects covered (only
 tables and databases, but not tablespaces and indexes).

The attached patch contributes:

- database_size(name)
- relation_size(text)
These are the well-known functions, tablespace-aware.

- pg_tablespace_size(oid)
- pg_database_size(oid)
- pg_relation_size(oid)
Tablespace-aware implementations, used by the upper functions.
pg_relation_size will report sizes of indexes as well.

- pg_size_pretty(bigint)
Formatting of sizes, to display '146MB' instead of '152885668'

Andreas Pflug
---
 contrib/dbsize/dbsize.c      | 391 ++++++++++++++++++++++++-----------
 contrib/dbsize/dbsize.sql.in |  16 ++
 2 files changed, 283 insertions(+), 124 deletions(-)

diff --git a/contrib/dbsize/dbsize.c b/contrib/dbsize/dbsize.c
index 9cf2c0f9369..8310e6f6cad 100644
--- a/contrib/dbsize/dbsize.c
+++ b/contrib/dbsize/dbsize.c
@@ -1,157 +1,285 @@
+/*
+ * dbsize.c
+ * object size functions
+ *
+ * Copyright (c) 2004, PostgreSQL Global Development Group
+ *
+ * Author: Andreas Pflug <pgadmin@pse-consulting.de>
+ *
+ * IDENTIFICATION
+ *	  $PostgreSQL: pgsql/contrib/dbsize/dbsize.c,v 1.13 2004/09/02 00:55:22 momjian Exp $
+ *
+ */
+
+
 #include "postgres.h"
 
 #include <sys/types.h>
 #include <sys/stat.h>
-#include <unistd.h>
 
 #include "access/heapam.h"
-#include "catalog/catalog.h"
-#include "catalog/catname.h"
+#include "storage/fd.h"
+#include "utils/syscache.h"
+#include "utils/builtins.h"
 #include "catalog/namespace.h"
 #include "catalog/pg_tablespace.h"
 #include "commands/dbcommands.h"
-#include "fmgr.h"
-#include "storage/fd.h"
-#include "utils/builtins.h"
+#include "miscadmin.h"
 
 
-static int64
-			get_tablespace_size(Oid dbid, Oid spcid, bool baddirOK);
+extern DLLIMPORT char *DataDir;
 
-static char *
-psnprintf(size_t len, const char *fmt,...)
-{
-	va_list		ap;
-	char	   *buf;
+Datum pg_tablespace_size(PG_FUNCTION_ARGS);
+Datum pg_database_size(PG_FUNCTION_ARGS);
+Datum pg_relation_size(PG_FUNCTION_ARGS);
+Datum pg_size_pretty(PG_FUNCTION_ARGS);
 
-	buf = palloc(len);
+Datum database_size(PG_FUNCTION_ARGS);
+Datum relation_size(PG_FUNCTION_ARGS);
 
-	va_start(ap, fmt);
-	vsnprintf(buf, len, fmt, ap);
-	va_end(ap);
+PG_FUNCTION_INFO_V1(pg_tablespace_size);
+PG_FUNCTION_INFO_V1(pg_database_size);
+PG_FUNCTION_INFO_V1(pg_relation_size);
+PG_FUNCTION_INFO_V1(pg_size_pretty);
 
-	return buf;
-}
+PG_FUNCTION_INFO_V1(database_size);
+PG_FUNCTION_INFO_V1(relation_size);
 
 
 
-/*
- * SQL function: database_size(name) returns bigint
- */
+static int64
+db_dir_size(char *path)
+{
+    int64 dirsize=0;
+    struct dirent *direntry;
+	DIR         *dirdesc;
+	char filename[MAXPGPATH];
 
-PG_FUNCTION_INFO_V1(database_size);
+	dirdesc=AllocateDir(path);
 
-Datum		database_size(PG_FUNCTION_ARGS);
+	if (!dirdesc)
+	    return 0;
 
-Datum
-database_size(PG_FUNCTION_ARGS)
-{
-	Name		dbname = PG_GETARG_NAME(0);
+	while ((direntry = readdir(dirdesc)) != 0)
+	{
+	    struct stat fst;
 
-	Oid			dbid;
-	int64		totalsize;
+	    if (!strcmp(direntry->d_name, ".") || !strcmp(direntry->d_name, ".."))
+		    continue;
 
-#ifdef SYMLINK
-	Relation	dbrel;
-	HeapScanDesc scan;
-	HeapTuple	tuple;
-#endif
+		snprintf(filename, MAXPGPATH, "%s/%s", path, direntry->d_name);
 
-	dbid = get_database_oid(NameStr(*dbname));
-	if (!OidIsValid(dbid))
-		ereport(ERROR,
-				(errcode(ERRCODE_UNDEFINED_DATABASE),
-			errmsg("database \"%s\" does not exist", NameStr(*dbname))));
+		if (stat(filename, &fst) < 0)
+			ereport(ERROR,
+					(errcode_for_file_access(),
+					 errmsg("could not stat \"%s\": %m", filename)));
+		dirsize += fst.st_size;
+	}
+
+	FreeDir(dirdesc);
+	return dirsize;
+}
+
+
+static int64
+calculate_database_size(Oid dbOid)
+{
+	int64 totalsize=0;
+	DIR         *dirdesc;
+    struct dirent *direntry;
+	char pathname[MAXPGPATH];
 
-#ifdef SYMLINK
+	snprintf(pathname, MAXPGPATH, "%s/global/%u", DataDir, (unsigned)dbOid);
+	totalsize += db_dir_size(pathname);
+	snprintf(pathname, MAXPGPATH, "%s/base/%u", DataDir, (unsigned)dbOid);
+	totalsize += db_dir_size(pathname);
 
-	dbrel = heap_openr(TableSpaceRelationName, AccessShareLock);
-	scan = heap_beginscan(dbrel, SnapshotNow, 0, (ScanKey) NULL);
+	snprintf(pathname, MAXPGPATH, "%s/pg_tblspc", DataDir);
+	dirdesc = AllocateDir(pathname);
 
-	totalsize = 0;
+	if (!dirdesc)
+	    ereport(ERROR,
+				(errcode_for_file_access(),
+				 errmsg("could not open tablespace directory: %m")));
 
-	while ((tuple = heap_getnext(scan, ForwardScanDirection)))
+	while ((direntry = readdir(dirdesc)) != 0)
 	{
-		Oid			spcid = HeapTupleGetOid(tuple);
+	    if (!strcmp(direntry->d_name, ".") || !strcmp(direntry->d_name, ".."))
+		    continue;
 
-		if (spcid != GLOBALTABLESPACE_OID)
-			totalsize += get_tablespace_size(dbid, spcid, true);
+		snprintf(pathname, MAXPGPATH, "%s/pg_tblspc/%s/%u", DataDir, direntry->d_name, (unsigned)dbOid);
+		totalsize += db_dir_size(pathname);
 	}
-	heap_endscan(scan);
-	heap_close(dbrel, AccessShareLock);
-#else
-	/* Same as always */
-	totalsize = get_tablespace_size(dbid, DEFAULTTABLESPACE_OID, false);
-#endif
-
-	/*
-	 * We need to keep in mind that we may not be called from the database
-	 * whose size we're reporting so, we need to look in every tablespace
-	 * to see if our database has data in there
-	 */
 
-	PG_RETURN_INT64(totalsize);
+	FreeDir(dirdesc);
+
+	if (!totalsize)
+	    ereport(ERROR,
+				(ERRCODE_UNDEFINED_DATABASE,
+				 errmsg("Database OID %u unknown.", (unsigned)dbOid)));
+
+	return totalsize;
 }
 
-static int64
-get_tablespace_size(Oid dbid, Oid spcid, bool baddirOK)
+/*
+ * calculate total size of tablespace
+ */
+Datum
+pg_tablespace_size(PG_FUNCTION_ARGS)
 {
-	char	   *dbpath;
-	DIR		   *dirdesc;
-	struct dirent *direntry;
-	int64		totalsize;
+    Oid tblspcOid = PG_GETARG_OID(0);
+
+	char tblspcPath[MAXPGPATH];
+	char pathname[MAXPGPATH];
+	int64		totalsize=0;
+	DIR         *dirdesc;
+    struct dirent *direntry;
 
-	dbpath = GetDatabasePath(dbid, spcid);
+	if (tblspcOid == DEFAULTTABLESPACE_OID)
+	    snprintf(tblspcPath, MAXPGPATH, "%s/base", DataDir);
+	else if (tblspcOid == GLOBALTABLESPACE_OID)
+	    snprintf(tblspcPath, MAXPGPATH, "%s/global", DataDir);
+	else
+		snprintf(tblspcPath, MAXPGPATH, "%s/pg_tblspc/%u", DataDir, (unsigned)tblspcOid);
+
+	dirdesc = AllocateDir(tblspcPath);
 
-	dirdesc = AllocateDir(dbpath);
 	if (!dirdesc)
+		ereport(ERROR,
+				(errcode_for_file_access(),
+				 errmsg("No such tablespace OID: %u: %m", (unsigned)tblspcOid)));
+
+	while ((direntry = readdir(dirdesc)) != 0)
 	{
-		if (baddirOK)
-			return 0;
-		else
+	    struct stat fst;
+
+	    if (!strcmp(direntry->d_name, ".") || !strcmp(direntry->d_name, ".."))
+		    continue;
+
+		snprintf(pathname, MAXPGPATH, "%s/%s", tblspcPath, direntry->d_name);
+		if (stat(pathname, &fst) < 0)
 			ereport(ERROR,
 					(errcode_for_file_access(),
-				 errmsg("could not open directory \"%s\": %m", dbpath)));
+					 errmsg("could not stat \"%s\": %m", pathname)));
+		totalsize += fst.st_size;
+
+		if (fst.st_mode & S_IFDIR)
+		    totalsize += db_dir_size(pathname);
 	}
-	totalsize = 0;
-	for (;;)
+
+	FreeDir(dirdesc);
+
+	PG_RETURN_INT64(totalsize);
+}
+
+
+/*
+ * calculate size of databases in all tablespaces
+ */
+Datum
+pg_database_size(PG_FUNCTION_ARGS)
+{
+    Oid dbOid = PG_GETARG_OID(0);
+
+	PG_RETURN_INT64(calculate_database_size(dbOid));
+}
+
+
+Datum
+database_size(PG_FUNCTION_ARGS)
+{
+	Name dbName = PG_GETARG_NAME(0);
+	Oid dbOid = get_database_oid(NameStr(*dbName));
+
+	if (!OidIsValid(dbOid))
+		ereport(ERROR,
+				(errcode(ERRCODE_UNDEFINED_DATABASE),
+			errmsg("database \"%s\" does not exist", NameStr(*dbName))));
+
+	PG_RETURN_INT64(calculate_database_size(dbOid));
+}
+
+static int64
+calculate_relation_size(Oid tblspcOid, Oid relnodeOid)
+{
+	int64		totalsize=0;
+	unsigned int segcount=0;
+	char dirpath[MAXPGPATH];
+	char pathname[MAXPGPATH];
+
+	if (tblspcOid == 0 || tblspcOid == DEFAULTTABLESPACE_OID)
+	    snprintf(dirpath, MAXPGPATH, "%s/base/%u", DataDir, (unsigned)MyDatabaseId);
+	else if (tblspcOid == GLOBALTABLESPACE_OID)
+	    snprintf(dirpath, MAXPGPATH, "%s/global", DataDir);
+	else
+	    snprintf(dirpath, MAXPGPATH, "%s/pg_tblspc/%u/%u", DataDir, (unsigned)tblspcOid, (unsigned)MyDatabaseId);
+
+	for (segcount = 0 ;; segcount++)
 	{
-		char	   *fullname;
-		struct stat statbuf;
+		struct stat fst;
+
+		if (segcount == 0)
+		    snprintf(pathname, MAXPGPATH, "%s/%u", dirpath, (unsigned) relnodeOid);
+		else
+		    snprintf(pathname, MAXPGPATH, "%s/%u.%u", dirpath, (unsigned) relnodeOid, segcount);
 
-		errno = 0;
-		direntry = readdir(dirdesc);
-		if (!direntry)
+		if (stat(pathname, &fst) < 0)
 		{
-			if (errno)
+			if (errno == ENOENT)
+				break;
+			else
 				ereport(ERROR,
 						(errcode_for_file_access(),
-						 errmsg("error reading directory: %m")));
-			else
-				break;
+						 errmsg("could not stat \"%s\": %m", pathname)));
 		}
-
-		fullname = psnprintf(strlen(dbpath) + 1 + strlen(direntry->d_name) + 1,
-							 "%s/%s", dbpath, direntry->d_name);
-		if (stat(fullname, &statbuf) == -1)
-			ereport(ERROR,
-					(errcode_for_file_access(),
-					 errmsg("could not stat \"%s\": %m", fullname)));
-		totalsize += statbuf.st_size;
-		pfree(fullname);
+		totalsize += fst.st_size;
 	}
 
-	FreeDir(dirdesc);
-	return (totalsize);
+	return totalsize;
 }
 
 /*
- * SQL function: relation_size(text) returns bigint
+ * calculate size of relation
  */
+Datum
+pg_relation_size(PG_FUNCTION_ARGS)
+{
+	Oid         relOid=PG_GETARG_OID(0);
 
-PG_FUNCTION_INFO_V1(relation_size);
+	HeapTuple   tuple;
+	Form_pg_class pg_class;
+	Oid			relnodeOid;
+	Oid         tblspcOid;
+    char        relkind;
+
+	tuple = SearchSysCache(RELOID, ObjectIdGetDatum(relOid), 0, 0, 0);
+	if (!HeapTupleIsValid(tuple))
+	    ereport(ERROR,
+				(ERRCODE_UNDEFINED_TABLE,
+				 errmsg("Relation OID %u does not exist", relOid)));
+
+	pg_class = (Form_pg_class) GETSTRUCT(tuple);
+	relnodeOid = pg_class->relfilenode;
+	tblspcOid = pg_class->reltablespace;
+	relkind = pg_class->relkind;
+
+	ReleaseSysCache(tuple);
+
+	switch(relkind)
+	{
+	    case RELKIND_INDEX:
+	    case RELKIND_RELATION:
+	    case RELKIND_TOASTVALUE:
+		    break;
+	    default:
+		    ereport(ERROR,
+					(ERRCODE_WRONG_OBJECT_TYPE,
+					 errmsg("Relation kind %d not supported", relkind)));
+	}
+
+	PG_RETURN_INT64(calculate_relation_size(tblspcOid, relnodeOid));
+}
 
-Datum		relation_size(PG_FUNCTION_ARGS);
 
 Datum
 relation_size(PG_FUNCTION_ARGS)
@@ -160,43 +288,58 @@ relation_size(PG_FUNCTION_ARGS)
 
 	RangeVar   *relrv;
 	Relation	relation;
-	Oid			relnode;
-	int64		totalsize;
-	unsigned int segcount;
+	Oid			relnodeOid;
+	Oid         tblspcOid;
 
 	relrv = makeRangeVarFromNameList(textToQualifiedNameList(relname,
 													   "relation_size"));
 	relation = heap_openrv(relrv, AccessShareLock);
 
-	relnode = relation->rd_rel->relfilenode;
+	tblspcOid  = relation->rd_rel->reltablespace;
+	relnodeOid = relation->rd_rel->relfilenode;
 
-	totalsize = 0;
-	segcount = 0;
-	for (;;)
-	{
-		char	   *fullname;
-		struct stat statbuf;
+	heap_close(relation, AccessShareLock);
 
-		if (segcount == 0)
-			fullname = psnprintf(25, "%u", (unsigned) relnode);
-		else
-			fullname = psnprintf(50, "%u.%u", (unsigned) relnode, segcount);
+	PG_RETURN_INT64(calculate_relation_size(tblspcOid, relnodeOid));
+}
 
-		if (stat(fullname, &statbuf) == -1)
+/*
+ * formatting with size units
+ */
+Datum
+pg_size_pretty(PG_FUNCTION_ARGS)
+{
+    int64 size=PG_GETARG_INT64(0);
+	char *result=palloc(50+VARHDRSZ);
+	int64 limit = 10*1024;
+	int64 mult=1;
+
+	if (size < limit*mult)
+	    snprintf(VARDATA(result), 50, INT64_FORMAT" bytes", size);
+    else
+	{
+		mult *= 1024;
+		if (size < limit*mult)
+		     snprintf(VARDATA(result), 50, INT64_FORMAT " kB", (size+mult/2) / mult);
+		else
 		{
-			if (errno == ENOENT)
-				break;
+			mult *= 1024;
+			if (size < limit*mult)
+			    snprintf(VARDATA(result), 50, INT64_FORMAT " MB", (size+mult/2) / mult);
 			else
-				ereport(ERROR,
-						(errcode_for_file_access(),
-						 errmsg("could not stat \"%s\": %m", fullname)));
+			{
+				mult *= 1024;
+				if (size < limit*mult)
+				    snprintf(VARDATA(result), 50, INT64_FORMAT " GB", (size+mult/2) / mult);
+				else
+				{
+				    mult *= 1024;
+				    snprintf(VARDATA(result), 50, INT64_FORMAT " TB", (size+mult/2) / mult);
+				}
+			}
 		}
-		totalsize += statbuf.st_size;
-		pfree(fullname);
-		segcount++;
 	}
+	VARATT_SIZEP(result) = strlen(VARDATA(result)) + VARHDRSZ;
 
-	heap_close(relation, AccessShareLock);
-
-	PG_RETURN_INT64(totalsize);
+	PG_RETURN_TEXT_P(result);
 }
diff --git a/contrib/dbsize/dbsize.sql.in b/contrib/dbsize/dbsize.sql.in
index b00a36fa200..a4ddc7e41fd 100644
--- a/contrib/dbsize/dbsize.sql.in
+++ b/contrib/dbsize/dbsize.sql.in
@@ -5,3 +5,19 @@ CREATE FUNCTION database_size (name) RETURNS bigint
 CREATE FUNCTION relation_size (text) RETURNS bigint
     AS 'MODULE_PATHNAME', 'relation_size'
     LANGUAGE C WITH (isstrict);
+
+CREATE FUNCTION pg_tablespace_size(oid) RETURNS bigint
+    AS 'MODULE_PATHNAME', 'pg_tablespace_size'
+    LANGUAGE C STABLE STRICT;
+
+CREATE FUNCTION pg_database_size(oid) RETURNS bigint
+    AS 'MODULE_PATHNAME', 'pg_database_size'
+    LANGUAGE C STABLE STRICT;
+
+CREATE FUNCTION pg_relation_size(oid) RETURNS bigint
+    AS 'MODULE_PATHNAME', 'pg_relation_size'
+    LANGUAGE C STABLE STRICT;
+
+CREATE FUNCTION pg_size_pretty(bigint) RETURNS text
+    AS 'MODULE_PATHNAME', 'pg_size_pretty'
+    LANGUAGE C STABLE STRICT;
-- 
GitLab