From fb7219266cee61f278dc8236c27ed63992cb0422 Mon Sep 17 00:00:00 2001
From: Christoph Berg <myon@debian.org>
Date: Fri, 14 Mar 2025 16:29:19 +0100
Subject: [PATCH v5] Add pg_tablespace_avail() functions

This exposes the f_avail value from statvfs() on tablespace directories
on the SQL level, allowing monitoring of free disk space from within the
server. On windows, GetDiskFreeSpaceEx() is used.

Permissions required match those from pg_tablespace_size().

In psql, include a new "Free" column in \db+ output.

Add test coverage for pg_tablespace_avail() and the previously not
covered pg_tablespace_size() function.
---
 doc/src/sgml/func/func-admin.sgml        |  21 +++++
 doc/src/sgml/ref/psql-ref.sgml           |   2 +-
 src/backend/utils/adt/dbsize.c           | 102 +++++++++++++++++++++++
 src/bin/psql/describe.c                  |  32 +++++--
 src/include/catalog/pg_proc.dat          |   8 ++
 src/test/regress/expected/tablespace.out |  21 +++++
 src/test/regress/sql/tablespace.sql      |  10 +++
 7 files changed, 189 insertions(+), 7 deletions(-)

diff --git a/doc/src/sgml/func/func-admin.sgml b/doc/src/sgml/func/func-admin.sgml
index 72038fc835f..959b0b673ab 100644
--- a/doc/src/sgml/func/func-admin.sgml
+++ b/doc/src/sgml/func/func-admin.sgml
@@ -1755,6 +1755,27 @@ postgres=# SELECT '0/0'::pg_lsn + pd.segment_number * ps.setting::int + :offset
        </para></entry>
       </row>
 
+      <row>
+       <entry role="func_table_entry"><para role="func_signature">
+        <indexterm>
+         <primary>pg_tablespace_avail</primary>
+        </indexterm>
+        <function>pg_tablespace_avail</function> ( <type>name</type> )
+        <returnvalue>bigint</returnvalue>
+       </para>
+       <para role="func_signature">
+        <function>pg_tablespace_avail</function> ( <type>oid</type> )
+        <returnvalue>bigint</returnvalue>
+       </para>
+       <para>
+        Returns the available disk space in the tablespace with the
+        specified name or OID. To use this function, you must
+        have <literal>CREATE</literal> privilege on the specified tablespace
+        or have privileges of the <literal>pg_read_all_stats</literal> role,
+        unless it is the default tablespace for the current database.
+       </para></entry>
+      </row>
+
       <row>
        <entry role="func_table_entry"><para role="func_signature">
         <indexterm>
diff --git a/doc/src/sgml/ref/psql-ref.sgml b/doc/src/sgml/ref/psql-ref.sgml
index 7c05afd4719..1ea67b3659f 100644
--- a/doc/src/sgml/ref/psql-ref.sgml
+++ b/doc/src/sgml/ref/psql-ref.sgml
@@ -1501,7 +1501,7 @@ SELECT $1 \parse stmt1
         If <literal>x</literal> is appended to the command name, the results
         are displayed in expanded mode.
         If <literal>+</literal> is appended to the command name, each tablespace
-        is listed with its associated options, on-disk size, permissions and
+        is listed with its associated options, on-disk size and free disk space, permissions and
         description.
         </para>
         </listitem>
diff --git a/src/backend/utils/adt/dbsize.c b/src/backend/utils/adt/dbsize.c
index cccc4a24c84..b395824ca3f 100644
--- a/src/backend/utils/adt/dbsize.c
+++ b/src/backend/utils/adt/dbsize.c
@@ -12,6 +12,12 @@
 #include "postgres.h"
 
 #include <sys/stat.h>
+#ifdef WIN32
+#include <fileapi.h>
+#include <errhandlingapi.h>
+#else
+#include <sys/statvfs.h>
+#endif
 
 #include "access/htup_details.h"
 #include "access/relation.h"
@@ -316,6 +322,102 @@ pg_tablespace_size_name(PG_FUNCTION_ARGS)
 }
 
 
+/*
+ * Return available disk space of tablespace. Returns -1 if the tablespace
+ * directory cannot be found.
+ */
+static int64
+calculate_tablespace_avail(Oid tblspcOid)
+{
+	char		tblspcPath[MAXPGPATH];
+	AclResult	aclresult;
+#ifdef WIN32
+	ULARGE_INTEGER lpFreeBytesAvailable;
+#else
+	struct statvfs fst;
+#endif
+
+	/*
+	 * User must have privileges of pg_read_all_stats or have CREATE privilege
+	 * for target tablespace, either explicitly granted or implicitly because
+	 * it is default for current database.
+	 */
+	if (tblspcOid != MyDatabaseTableSpace &&
+		!has_privs_of_role(GetUserId(), ROLE_PG_READ_ALL_STATS))
+	{
+		aclresult = object_aclcheck(TableSpaceRelationId, tblspcOid, GetUserId(), ACL_CREATE);
+		if (aclresult != ACLCHECK_OK)
+			aclcheck_error(aclresult, OBJECT_TABLESPACE,
+						   get_tablespace_name(tblspcOid));
+	}
+
+	if (tblspcOid == DEFAULTTABLESPACE_OID)
+		snprintf(tblspcPath, MAXPGPATH, "base");
+	else if (tblspcOid == GLOBALTABLESPACE_OID)
+		snprintf(tblspcPath, MAXPGPATH, "global");
+	else
+		snprintf(tblspcPath, MAXPGPATH, "%s/%u/%s", PG_TBLSPC_DIR, tblspcOid,
+				 TABLESPACE_VERSION_DIRECTORY);
+
+#ifdef WIN32
+	if (! GetDiskFreeSpaceEx(tblspcPath, &lpFreeBytesAvailable, NULL, NULL))
+		elog(ERROR, "GetDiskFreeSpaceEx failed: error code %lu", GetLastError());
+
+	return lpFreeBytesAvailable.QuadPart; /* ULONGLONG part of ULARGE_INTEGER */
+#else
+	if (statvfs(tblspcPath, &fst) < 0)
+	{
+		if (errno == ENOENT)
+			return -1;
+		else
+			ereport(ERROR,
+					(errcode_for_file_access(),
+					 errmsg("could not statvfs directory \"%s\": %m", tblspcPath)));
+	}
+
+	return fst.f_bavail * fst.f_frsize; /* available blocks times fragment size */
+#endif
+}
+
+Datum
+pg_tablespace_avail_oid(PG_FUNCTION_ARGS)
+{
+	Oid			tblspcOid = PG_GETARG_OID(0);
+	int64		avail;
+
+	/*
+	 * Not needed for correctness, but avoid non-user-facing error message
+	 * later if the tablespace doesn't exist.
+	 */
+	if (!SearchSysCacheExists1(TABLESPACEOID, ObjectIdGetDatum(tblspcOid)))
+		ereport(ERROR,
+				errcode(ERRCODE_UNDEFINED_OBJECT),
+				errmsg("tablespace with OID %u does not exist", tblspcOid));
+
+	avail = calculate_tablespace_avail(tblspcOid);
+
+	if (avail < 0)
+		PG_RETURN_NULL();
+
+	PG_RETURN_INT64(avail);
+}
+
+Datum
+pg_tablespace_avail_name(PG_FUNCTION_ARGS)
+{
+	Name		tblspcName = PG_GETARG_NAME(0);
+	Oid			tblspcOid = get_tablespace_oid(NameStr(*tblspcName), false);
+	int64		avail;
+
+	avail = calculate_tablespace_avail(tblspcOid);
+
+	if (avail < 0)
+		PG_RETURN_NULL();
+
+	PG_RETURN_INT64(avail);
+}
+
+
 /*
  * calculate size of (one fork of) a relation
  *
diff --git a/src/bin/psql/describe.c b/src/bin/psql/describe.c
index e1449654f96..36486417a48 100644
--- a/src/bin/psql/describe.c
+++ b/src/bin/psql/describe.c
@@ -234,7 +234,7 @@ describeTablespaces(const char *pattern, bool verbose)
 	appendPQExpBuffer(&buf,
 					  "SELECT spcname AS \"%s\",\n"
 					  "  pg_catalog.pg_get_userbyid(spcowner) AS \"%s\",\n"
-					  "  pg_catalog.pg_tablespace_location(oid) AS \"%s\"",
+					  "  pg_catalog.pg_tablespace_location(tblspc.oid) AS \"%s\"",
 					  gettext_noop("Name"),
 					  gettext_noop("Owner"),
 					  gettext_noop("Location"));
@@ -245,15 +245,34 @@ describeTablespaces(const char *pattern, bool verbose)
 		printACLColumn(&buf, "spcacl");
 		appendPQExpBuffer(&buf,
 						  ",\n  spcoptions AS \"%s\""
-						  ",\n  pg_catalog.pg_size_pretty(pg_catalog.pg_tablespace_size(oid)) AS \"%s\""
-						  ",\n  pg_catalog.shobj_description(oid, 'pg_tablespace') AS \"%s\"",
+						  ",\n  CASE WHEN dbsub.dattablespace OPERATOR(pg_catalog.=) tblspc.oid OR\n"
+						  "               pg_catalog.has_tablespace_privilege(tblspc.oid, 'CREATE') OR\n"
+						  "               pg_catalog.pg_has_role('pg_read_all_stats', 'USAGE')\n"
+						  "       THEN pg_catalog.pg_size_pretty(pg_catalog.pg_tablespace_size(tblspc.oid))\n"
+						  "       ELSE 'No Access'"
+						  "  END as \"%s\"",
 						  gettext_noop("Options"),
-						  gettext_noop("Size"),
+						  gettext_noop("Size"));
+		if (pset.sversion >= 190000)
+			appendPQExpBuffer(&buf,
+							  ",\n  CASE WHEN dbsub.dattablespace OPERATOR(pg_catalog.=) tblspc.oid OR\n"
+							  "               pg_catalog.has_tablespace_privilege(tblspc.oid, 'CREATE') OR\n"
+							  "               pg_catalog.pg_has_role('pg_read_all_stats', 'USAGE')\n"
+							  "       THEN pg_catalog.pg_size_pretty(pg_catalog.pg_tablespace_avail(tblspc.oid))\n"
+							  "       ELSE 'No Access'"
+							  "  END as \"%s\"",
+							  gettext_noop("Free"));
+		appendPQExpBuffer(&buf,
+						  ",\n  pg_catalog.shobj_description(tblspc.oid, 'pg_tablespace') AS \"%s\"",
 						  gettext_noop("Description"));
 	}
 
 	appendPQExpBufferStr(&buf,
-						 "\nFROM pg_catalog.pg_tablespace\n");
+						 "\nFROM pg_catalog.pg_tablespace tblspc\n");
+	if (verbose)
+		appendPQExpBufferStr(&buf,
+							 "CROSS JOIN (SELECT dattablespace FROM pg_catalog.pg_database db\n"
+							 "  wHERE db.datname OPERATOR(pg_catalog.=) pg_catalog.current_database()) dbsub\n");
 
 	if (!validateSQLNamePattern(&buf, pattern, false, false,
 								NULL, "spcname", NULL,
@@ -1008,7 +1027,8 @@ listAllDbs(const char *pattern, bool verbose)
 	printACLColumn(&buf, "d.datacl");
 	if (verbose)
 		appendPQExpBuffer(&buf,
-						  ",\n  CASE WHEN pg_catalog.has_database_privilege(d.datname, 'CONNECT')\n"
+						  ",\n  CASE WHEN pg_catalog.has_database_privilege(d.datname, 'CONNECT') OR\n"
+						  "               pg_catalog.pg_has_role('pg_read_all_stats', 'USAGE')\n"
 						  "       THEN pg_catalog.pg_size_pretty(pg_catalog.pg_database_size(d.datname))\n"
 						  "       ELSE 'No Access'\n"
 						  "  END as \"%s\""
diff --git a/src/include/catalog/pg_proc.dat b/src/include/catalog/pg_proc.dat
index fa9ae79082b..78c03ea6412 100644
--- a/src/include/catalog/pg_proc.dat
+++ b/src/include/catalog/pg_proc.dat
@@ -7866,6 +7866,14 @@
   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 => '6015',
+  descr => 'disk stats for the specified tablespace',
+  proname => 'pg_tablespace_avail', provolatile => 'v', prorettype => 'int8',
+  proargtypes => 'oid', prosrc => 'pg_tablespace_avail_oid' },
+{ oid => '6016',
+  descr => 'disk stats for the specified tablespace',
+  proname => 'pg_tablespace_avail', provolatile => 'v', prorettype => 'int8',
+  proargtypes => 'name', prosrc => 'pg_tablespace_avail_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' },
diff --git a/src/test/regress/expected/tablespace.out b/src/test/regress/expected/tablespace.out
index f0dd25cdf0c..12a78c77e05 100644
--- a/src/test/regress/expected/tablespace.out
+++ b/src/test/regress/expected/tablespace.out
@@ -20,6 +20,27 @@ SELECT spcoptions FROM pg_tablespace WHERE spcname = 'regress_tblspacewith';
  {random_page_cost=3.0}
 (1 row)
 
+-- check size functions
+SELECT pg_tablespace_size('pg_default') BETWEEN 1_000_000 and 10_000_000_000, -- rough sanity check
+       pg_tablespace_size('pg_global') BETWEEN 100_000 and 10_000_000,
+       pg_tablespace_size('regress_tblspacewith'); -- empty
+ ?column? | ?column? | pg_tablespace_size 
+----------+----------+--------------------
+ t        | t        |                  0
+(1 row)
+
+SELECT pg_tablespace_size('missing');
+ERROR:  tablespace "missing" does not exist
+SELECT pg_tablespace_avail('pg_default') > 1_000_000,
+       pg_tablespace_avail('pg_global') > 1_000_000,
+       pg_tablespace_avail('regress_tblspacewith') > 1_000_000;
+ ?column? | ?column? | ?column? 
+----------+----------+----------
+ t        | t        | t
+(1 row)
+
+SELECT pg_tablespace_avail('missing');
+ERROR:  tablespace "missing" does not exist
 -- drop the tablespace so we can re-use the location
 DROP TABLESPACE regress_tblspacewith;
 -- This returns a relative path as of an effect of allow_in_place_tablespaces,
diff --git a/src/test/regress/sql/tablespace.sql b/src/test/regress/sql/tablespace.sql
index c43a59e5957..91152335459 100644
--- a/src/test/regress/sql/tablespace.sql
+++ b/src/test/regress/sql/tablespace.sql
@@ -17,6 +17,16 @@ CREATE TABLESPACE regress_tblspacewith LOCATION '' WITH (random_page_cost = 3.0)
 -- check to see the parameter was used
 SELECT spcoptions FROM pg_tablespace WHERE spcname = 'regress_tblspacewith';
 
+-- check size functions
+SELECT pg_tablespace_size('pg_default') BETWEEN 1_000_000 and 10_000_000_000, -- rough sanity check
+       pg_tablespace_size('pg_global') BETWEEN 100_000 and 10_000_000,
+       pg_tablespace_size('regress_tblspacewith'); -- empty
+SELECT pg_tablespace_size('missing');
+SELECT pg_tablespace_avail('pg_default') > 1_000_000,
+       pg_tablespace_avail('pg_global') > 1_000_000,
+       pg_tablespace_avail('regress_tblspacewith') > 1_000_000;
+SELECT pg_tablespace_avail('missing');
+
 -- drop the tablespace so we can re-use the location
 DROP TABLESPACE regress_tblspacewith;
 
-- 
2.53.0

