Use a lookup table for units in pg_size_pretty and pg_size_bytes

We've grown 2 versions of pg_size_pretty over the years, one for BIGINT
and one for NUMERIC.  Both should output the same, but keeping them in
sync is harder than needed due to neither function sharing a source of
truth about which units to use and how to transition to the next largest
unit.

Here we add a static array which defines the units that we recognize and
have both pg_size_pretty and pg_size_pretty_numeric use it.  This will
make adding any units in the future a very simple task.

The table contains all information required to allow us to also modify
pg_size_bytes to use the lookup table, so adjust that too.

There are no behavioral changes here.

Author: David Rowley
Reviewed-by: Dean Rasheed, Tom Lane, David Christensen
Discussion: https://postgr.es/m/CAApHDvru1F7qsEVL-iOHeezJ+5WVxXnyD_Jo9nht+Eh85ekK-Q@mail.gmail.com

diff --git a/src/backend/utils/adt/dbsize.c b/src/backend/utils/adt/dbsize.c
index 9de2ed09d994a2e29750f89dccb7c5167e101d0b..6c381b02b7ec68b46a7b87e81dc82db086f186d7 100644 (file)
--- a/src/backend/utils/adt/dbsize.c
+++ b/src/backend/utils/adt/dbsize.c
@@ -34,6 +34,27 @@
 /* Divide by two and round away from zero */
 #define half_rounded(x)   (((x) + ((x) < 0 ? -1 : 1)) / 2)
 
+/* Units used in pg_size_pretty functions.  All units must be powers of 2 */
+struct size_pretty_unit
+{
+   const char *name;           /* bytes, kB, MB, GB etc */
+   uint32      limit;          /* upper limit, prior to half rounding after
+                                * converting to this unit. */
+   bool        round;          /* do half rounding for this unit */
+   uint8       unitbits;       /* (1 << unitbits) bytes to make 1 of this
+                                * unit */
+};
+
+/* When adding units here also update 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},
+   {"MB", 20 * 1024 - 1, true, 20},
+   {"GB", 20 * 1024 - 1, true, 30},
+   {"TB", 20 * 1024 - 1, true, 40},
+   {NULL, 0, false, 0}
+};
+
 /* Return physical size of directory contents, or 0 if dir doesn't exist */
 static int64
 db_dir_size(const char *path)
@@ -535,41 +556,34 @@ pg_size_pretty(PG_FUNCTION_ARGS)
 {
    int64       size = PG_GETARG_INT64(0);
    char        buf[64];
-   int64       limit = 10 * 1024;
-   int64       limit2 = limit * 2 - 1;
+   const struct size_pretty_unit *unit;
 
-   if (Abs(size) < limit)
-       snprintf(buf, sizeof(buf), INT64_FORMAT " bytes", size);
-   else
+   for (unit = size_pretty_units; unit->name != NULL; unit++)
    {
-       /*
-        * We use divide instead of bit shifting so that behavior matches for
-        * both positive and negative size values.
-        */
-       size /= (1 << 9);       /* keep one extra bit for rounding */
-       if (Abs(size) < limit2)
-           snprintf(buf, sizeof(buf), INT64_FORMAT " kB",
-                    half_rounded(size));
-       else
+       uint8       bits;
+
+       /* use this unit if there are no more units or we're below the limit */
+       if (unit[1].name == NULL || Abs(size) < unit->limit)
        {
-           size /= (1 << 10);
-           if (Abs(size) < limit2)
-               snprintf(buf, sizeof(buf), INT64_FORMAT " MB",
-                        half_rounded(size));
-           else
-           {
-               size /= (1 << 10);
-               if (Abs(size) < limit2)
-                   snprintf(buf, sizeof(buf), INT64_FORMAT " GB",
-                            half_rounded(size));
-               else
-               {
-                   size /= (1 << 10);
-                   snprintf(buf, sizeof(buf), INT64_FORMAT " TB",
-                            half_rounded(size));
-               }
-           }
+           if (unit->round)
+               size = half_rounded(size);
+
+           snprintf(buf, sizeof(buf), INT64_FORMAT " %s", size, unit->name);
+           break;
        }
+
+       /*
+        * Determine the number of bits to use to build the divisor.  We may
+        * need to use 1 bit less than the difference between this and the
+        * next unit if the next unit uses half rounding.  Or we may need to
+        * shift an extra bit if this unit uses half rounding and the next one
+        * does not.  We use division rather than shifting right by this
+        * number of bits to ensure positive and negative values are rounded
+        * in the same way.
+        */
+       bits = (unit[1].unitbits - unit->unitbits - (unit[1].round == true)
+               + (unit->round == true));
+       size /= ((int64) 1) << bits;
    }
 
    PG_RETURN_TEXT_P(cstring_to_text(buf));
@@ -640,57 +654,35 @@ Datum
 pg_size_pretty_numeric(PG_FUNCTION_ARGS)
 {
    Numeric     size = PG_GETARG_NUMERIC(0);
-   Numeric     limit,
-               limit2;
-   char       *result;
-
-   limit = int64_to_numeric(10 * 1024);
-   limit2 = int64_to_numeric(10 * 1024 * 2 - 1);
+   char       *result = NULL;
+   const struct size_pretty_unit *unit;
 
-   if (numeric_is_less(numeric_absolute(size), limit))
+   for (unit = size_pretty_units; unit->name != NULL; unit++)
    {
-       result = psprintf("%s bytes", numeric_to_cstring(size));
-   }
-   else
-   {
-       /* keep one extra bit for rounding */
-       /* size /= (1 << 9) */
-       size = numeric_truncated_divide(size, 1 << 9);
+       unsigned int shiftby;
 
-       if (numeric_is_less(numeric_absolute(size), limit2))
+       /* use this unit if there are no more units or we're below the limit */
+       if (unit[1].name == NULL ||
+           numeric_is_less(numeric_absolute(size),
+                           int64_to_numeric(unit->limit)))
        {
-           size = numeric_half_rounded(size);
-           result = psprintf("%s kB", numeric_to_cstring(size));
-       }
-       else
-       {
-           /* size /= (1 << 10) */
-           size = numeric_truncated_divide(size, 1 << 10);
-
-           if (numeric_is_less(numeric_absolute(size), limit2))
-           {
+           if (unit->round)
                size = numeric_half_rounded(size);
-               result = psprintf("%s MB", numeric_to_cstring(size));
-           }
-           else
-           {
-               /* size /= (1 << 10) */
-               size = numeric_truncated_divide(size, 1 << 10);
-
-               if (numeric_is_less(numeric_absolute(size), limit2))
-               {
-                   size = numeric_half_rounded(size);
-                   result = psprintf("%s GB", numeric_to_cstring(size));
-               }
-               else
-               {
-                   /* size /= (1 << 10) */
-                   size = numeric_truncated_divide(size, 1 << 10);
-                   size = numeric_half_rounded(size);
-                   result = psprintf("%s TB", numeric_to_cstring(size));
-               }
-           }
+
+           result = psprintf("%s %s", numeric_to_cstring(size), unit->name);
+           break;
        }
+
+       /*
+        * Determine the number of bits to use to build the divisor.  We may
+        * need to use 1 bit less than the difference between this and the
+        * next unit if the next unit uses half rounding.  Or we may need to
+        * shift an extra bit if this unit uses half rounding and the next one
+        * does not.
+        */
+       shiftby = (unit[1].unitbits - unit->unitbits - (unit[1].round == true)
+                  + (unit->round == true));
+       size = numeric_truncated_divide(size, ((int64) 1) << shiftby);
    }
 
    PG_RETURN_TEXT_P(cstring_to_text(result));
@@ -791,6 +783,7 @@ pg_size_bytes(PG_FUNCTION_ARGS)
    /* Handle possible unit */
    if (*strptr != '\0')
    {
+       const struct size_pretty_unit *unit;
        int64       multiplier = 0;
 
        /* Trim any trailing whitespace */
@@ -802,21 +795,18 @@ pg_size_bytes(PG_FUNCTION_ARGS)
        endptr++;
        *endptr = '\0';
 
-       /* Parse the unit case-insensitively */
-       if (pg_strcasecmp(strptr, "bytes") == 0)
-           multiplier = (int64) 1;
-       else if (pg_strcasecmp(strptr, "kb") == 0)
-           multiplier = (int64) 1024;
-       else if (pg_strcasecmp(strptr, "mb") == 0)
-           multiplier = ((int64) 1024) * 1024;
-
-       else if (pg_strcasecmp(strptr, "gb") == 0)
-           multiplier = ((int64) 1024) * 1024 * 1024;
-
-       else if (pg_strcasecmp(strptr, "tb") == 0)
-           multiplier = ((int64) 1024) * 1024 * 1024 * 1024;
+       for (unit = size_pretty_units; unit->name != NULL; unit++)
+       {
+           /* Parse the unit case-insensitively */
+           if (pg_strcasecmp(strptr, unit->name) == 0)
+           {
+               multiplier = ((int64) 1) << unit->unitbits;
+               break;
+           }
+       }
 
-       else
+       /* Verify we found a valid unit in the loop above */
+       if (unit->name == NULL)
            ereport(ERROR,
                    (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
                     errmsg("invalid size: \"%s\"", text_to_cstring(arg)),