Add support for unit "B" to pg_size_bytes()

This makes it consistent with the units support in GUC.

Reviewed-by: David Rowley <dgrowleyml@gmail.com>
Reviewed-by: Dean Rasheed <dean.a.rasheed@gmail.com>
Discussion: https://www.postgresql.org/message-id/flat/0106914a-9eb5-22be-40d8-652cc88c827d%40enterprisedb.com

3 files changed, 38 insertions, 12 deletions

diff --git a/src/backend/utils/adt/dbsize.c b/src/backend/utils/adt/dbsize.c
index dbd404101f..e5c0f1c45b 100644
--- a/src/backend/utils/adt/dbsize.c
+++ b/src/backend/utils/adt/dbsize.c
@@ -46,7 +46,7 @@ struct size_pretty_unit
 								 * unit */
 };
 
-/* When adding units here also update the error message in pg_size_bytes */
+/* When adding units here also update the docs and the error message in pg_size_bytes */
 static const struct size_pretty_unit size_pretty_units[] = {
 	{"bytes", 10 * 1024, false, 0},
 	{"kB", 20 * 1024 - 1, true, 10},
@@ -57,6 +57,19 @@ static const struct size_pretty_unit size_pretty_units[] = {
 	{NULL, 0, false, 0}
 };
 
+/* Additional unit aliases accepted by pg_size_bytes */
+struct size_bytes_unit_alias
+{
+	const char *alias;
+	int			unit_index;		/* corresponding size_pretty_units element */
+};
+
+/* When adding units here also update the docs and the error message in pg_size_bytes */
+static const struct size_bytes_unit_alias size_bytes_aliases[] = {
+	{"B", 0},
+	{NULL}
+};
+
 /* Return physical size of directory contents, or 0 if dir doesn't exist */
 static int64
 db_dir_size(const char *path)
@@ -801,9 +814,19 @@ pg_size_bytes(PG_FUNCTION_ARGS)
 		{
 			/* Parse the unit case-insensitively */
 			if (pg_strcasecmp(strptr, unit->name) == 0)
-			{
-				multiplier = ((int64) 1) << unit->unitbits;
 				break;
+		}
+
+		/* If not found, look in table of aliases */
+		if (unit->name == NULL)
+		{
+			for (const struct size_bytes_unit_alias *a = size_bytes_aliases; a->alias != NULL; a++)
+			{
+				if (pg_strcasecmp(strptr, a->alias) == 0)
+				{
+					unit = &size_pretty_units[a->unit_index];
+					break;
+				}
 			}
 		}
 
@@ -813,7 +836,9 @@ pg_size_bytes(PG_FUNCTION_ARGS)
 					(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
 					 errmsg("invalid size: \"%s\"", text_to_cstring(arg)),
 					 errdetail("Invalid size unit: \"%s\".", strptr),
-					 errhint("Valid units are \"bytes\", \"kB\", \"MB\", \"GB\", \"TB\", and \"PB\".")));
+					 errhint("Valid units are \"bytes\", \"B\", \"kB\", \"MB\", \"GB\", \"TB\", and \"PB\".")));
+
+		multiplier = ((int64) 1) << unit->unitbits;
 
 		if (multiplier > 1)
 		{
diff --git a/src/test/regress/expected/dbsize.out b/src/test/regress/expected/dbsize.out
index d8d6686b5f..f1121a87aa 100644
--- a/src/test/regress/expected/dbsize.out
+++ b/src/test/regress/expected/dbsize.out
@@ -81,12 +81,13 @@ SELECT size, pg_size_pretty(size), pg_size_pretty(-1 * size) FROM
 
 -- pg_size_bytes() tests
 SELECT size, pg_size_bytes(size) FROM
-    (VALUES ('1'), ('123bytes'), ('1kB'), ('1MB'), (' 1 GB'), ('1.5 GB '),
+    (VALUES ('1'), ('123bytes'), ('256 B'), ('1kB'), ('1MB'), (' 1 GB'), ('1.5 GB '),
             ('1TB'), ('3000 TB'), ('1e6 MB'), ('99 PB')) x(size);
    size   |   pg_size_bytes    
 ----------+--------------------
  1        |                  1
  123bytes |                123
+ 256 B    |                256
  1kB      |               1024
  1MB      |            1048576
   1 GB    |         1073741824
@@ -95,7 +96,7 @@ SELECT size, pg_size_bytes(size) FROM
  3000 TB  |   3298534883328000
  1e6 MB   |      1048576000000
  99 PB    | 111464090777419776
-(10 rows)
+(11 rows)
 
 -- case-insensitive units are supported
 SELECT size, pg_size_bytes(size) FROM
@@ -153,15 +154,15 @@ SELECT size, pg_size_bytes(size) FROM
 SELECT pg_size_bytes('1 AB');
 ERROR:  invalid size: "1 AB"
 DETAIL:  Invalid size unit: "AB".
-HINT:  Valid units are "bytes", "kB", "MB", "GB", "TB", and "PB".
+HINT:  Valid units are "bytes", "B", "kB", "MB", "GB", "TB", and "PB".
 SELECT pg_size_bytes('1 AB A');
 ERROR:  invalid size: "1 AB A"
 DETAIL:  Invalid size unit: "AB A".
-HINT:  Valid units are "bytes", "kB", "MB", "GB", "TB", and "PB".
+HINT:  Valid units are "bytes", "B", "kB", "MB", "GB", "TB", and "PB".
 SELECT pg_size_bytes('1 AB A    ');
 ERROR:  invalid size: "1 AB A    "
 DETAIL:  Invalid size unit: "AB A".
-HINT:  Valid units are "bytes", "kB", "MB", "GB", "TB", and "PB".
+HINT:  Valid units are "bytes", "B", "kB", "MB", "GB", "TB", and "PB".
 SELECT pg_size_bytes('9223372036854775807.9');
 ERROR:  bigint out of range
 SELECT pg_size_bytes('1e100');
@@ -171,7 +172,7 @@ ERROR:  value overflows numeric format
 SELECT pg_size_bytes('1 byte');  -- the singular "byte" is not supported
 ERROR:  invalid size: "1 byte"
 DETAIL:  Invalid size unit: "byte".
-HINT:  Valid units are "bytes", "kB", "MB", "GB", "TB", and "PB".
+HINT:  Valid units are "bytes", "B", "kB", "MB", "GB", "TB", and "PB".
 SELECT pg_size_bytes('');
 ERROR:  invalid size: ""
 SELECT pg_size_bytes('kb');
@@ -189,6 +190,6 @@ ERROR:  invalid size: ".+912"
 SELECT pg_size_bytes('+912+ kB');
 ERROR:  invalid size: "+912+ kB"
 DETAIL:  Invalid size unit: "+ kB".
-HINT:  Valid units are "bytes", "kB", "MB", "GB", "TB", and "PB".
+HINT:  Valid units are "bytes", "B", "kB", "MB", "GB", "TB", and "PB".
 SELECT pg_size_bytes('++123 kB');
 ERROR:  invalid size: "++123 kB"
diff --git a/src/test/regress/sql/dbsize.sql b/src/test/regress/sql/dbsize.sql
index 7df865271b..b34cf33385 100644
--- a/src/test/regress/sql/dbsize.sql
+++ b/src/test/regress/sql/dbsize.sql
@@ -29,7 +29,7 @@ SELECT size, pg_size_pretty(size), pg_size_pretty(-1 * size) FROM
 
 -- pg_size_bytes() tests
 SELECT size, pg_size_bytes(size) FROM
-    (VALUES ('1'), ('123bytes'), ('1kB'), ('1MB'), (' 1 GB'), ('1.5 GB '),
+    (VALUES ('1'), ('123bytes'), ('256 B'), ('1kB'), ('1MB'), (' 1 GB'), ('1.5 GB '),
             ('1TB'), ('3000 TB'), ('1e6 MB'), ('99 PB')) x(size);
 
 -- case-insensitive units are supported