Merge pull request #452 from diffix/piotr/normalize-for-metabase-date_trunc

Allow `cast(extract ...) as integer`

Author: pdobacz

pg_diffix/query/allowed_objects.h

@@ -2,6 +2,7 @@
 #define PG_DIFFIX_ALLOWED_OBJECTS_H
 
 #include "nodes/bitmapset.h"
+#include "nodes/primnodes.h"
 
 /*
  * Returns whether the OID points to a function (or operator) allowed in defining buckets.
@@ -17,7 +18,7 @@ extern int primary_arg_index(Oid funcoid);
 /*
  * Returns whether the OID points to a cast allowed in defining buckets.
  */
-extern bool is_allowed_cast(Oid funcoid);
+extern bool is_allowed_cast(const FuncExpr *func_expr);
 
 /*
  * Returns whether the OID points to a UDF being a implicit_range function, e.g. `ceil_by(x, 2.0)`,

src/node_funcs.c

@@ -20,7 +20,7 @@ Node *unwrap_cast(Node *node)
   if (IsA(node, FuncExpr))
   {
     FuncExpr *func_expr = (FuncExpr *)node;
-    if (is_allowed_cast(func_expr->funcid))
+    if (is_allowed_cast(func_expr))
     {
       Assert(list_length(func_expr->args) == 1); /* All allowed casts require exactly one argument. */
       return unwrap_cast(linitial(func_expr->args));

src/query/allowed_objects.c

@@ -1,6 +1,8 @@
 #include "postgres.h"
 
 #include "access/sysattr.h"
+#include "catalog/pg_type.h"
+#include "utils/builtins.h"
 #include "utils/fmgroids.h"
 #include "utils/fmgrtab.h"
 #include "utils/lsyscache.h"
@@ -13,7 +15,7 @@
 static const char *const g_allowed_casts[] = {
     "i2tod", "i2tof", "i2toi4", "i4toi2", "i4tod", "i4tof", "i8tod", "i8tof", "int48", "int84",
     "ftod", "dtof",
-    "int4_numeric", "float4_numeric", "float8_numeric",
+    "int2_numeric", "int4_numeric", "int8_numeric", "float4_numeric", "float8_numeric",
     "numeric_float4", "numeric_float8",
     "date_timestamptz",
     /**/
@@ -39,7 +41,9 @@ static const FunctionByName g_allowed_builtins[] = {
     (FunctionByName){.name = "dtoi2", .primary_arg = 0},
     (FunctionByName){.name = "dtoi4", .primary_arg = 0},
     (FunctionByName){.name = "dtoi8", .primary_arg = 0},
+    (FunctionByName){.name = "numeric_int2", .primary_arg = 0},
     (FunctionByName){.name = "numeric_int4", .primary_arg = 0},
+    (FunctionByName){.name = "numeric_int8", .primary_arg = 0},
     /* substring */
     (FunctionByName){.name = "text_substr", .primary_arg = 0},
     (FunctionByName){.name = "text_substr_no_len", .primary_arg = 0},
@@ -124,6 +128,16 @@ static AllowedCols g_pg_catalog_allowed_cols[] = {
     /**/
 };
 
+static const char *const g_decimal_integer_casts[] = {
+    "numeric_int2", "numeric_int4", "numeric_int8", "dtoi2", "dtoi4", "dtoi8", "ftoi2", "ftoi4", "ftoi8",
+    /**/
+};
+
+static const char *const g_extract_functions[] = {
+    "extract_date", "extract_timestamp", "extract_timestamptz", "timestamp_part", "timestamptz_part",
+    /**/
+};
+
 static void prepare_pg_catalog_allowed(Oid relation_oid, AllowedCols *allowed_cols)
 {
   MemoryContext old_context = MemoryContextSwitchTo(TopMemoryContext);
@@ -166,6 +180,16 @@ static const FmgrBuiltin *fmgr_isbuiltin(Oid id)
   return &fmgr_builtins[index];
 }
 
+static bool is_member_of(const char *s, const char *const array[], int length)
+{
+  for (int i = 0; i < length; i++)
+  {
+    if (strcmp(array[i], s) == 0)
+      return true;
+  }
+  return false;
+}
+
 static bool is_func_member_of(Oid funcoid, const FunctionByName func_array[], int length)
 {
   const FmgrBuiltin *fmgr_builtin = fmgr_isbuiltin(funcoid);
@@ -185,13 +209,7 @@ static bool is_funcname_member_of(Oid funcoid, const char *const name_array[], i
 {
   const FmgrBuiltin *fmgr_builtin = fmgr_isbuiltin(funcoid);
   if (fmgr_builtin != NULL)
-  {
-    for (int i = 0; i < length; i++)
-    {
-      if (strcmp(name_array[i], fmgr_builtin->funcName) == 0)
-        return true;
-    }
-  }
+    return is_member_of(fmgr_builtin->funcName, name_array, length);
 
   return false;
 }
@@ -224,9 +242,25 @@ int primary_arg_index(Oid funcoid)
   FAILWITH("Cannot identify the primary argument position for funcid %u.", funcoid);
 }
 
-bool is_allowed_cast(Oid funcoid)
+bool is_allowed_cast(const FuncExpr *func_expr)
 {
-  return is_funcname_member_of(funcoid, g_allowed_casts, ARRAY_LENGTH(g_allowed_casts));
+  if (is_funcname_member_of(func_expr->funcid, g_allowed_casts, ARRAY_LENGTH(g_allowed_casts)))
+  {
+    return true;
+  }
+  else if (is_funcname_member_of(func_expr->funcid, g_decimal_integer_casts, ARRAY_LENGTH(g_decimal_integer_casts)))
+  {
+    /* Handle cases like `cast(extract(minute from ...) as integer)`. */
+    Node *cast_arg = linitial(func_expr->args);
+    if (IsA(cast_arg, FuncExpr))
+    {
+      FuncExpr *cast_arg_expr = (FuncExpr *)cast_arg;
+      if (is_funcname_member_of(cast_arg_expr->funcid, g_extract_functions, ARRAY_LENGTH(g_extract_functions)) ||
+          cast_arg_expr->funcid == F_DATE_PART_TEXT_DATE)
+        return true;
+    }
+  }
+  return false;
 }
 
 bool is_implicit_range_udf_untrusted(Oid funcoid)
@@ -279,13 +313,10 @@ bool is_allowed_pg_catalog_rte(Oid relation_oid, const Bitmapset *selected_cols)
   char *rel_name = get_rel_name(relation_oid);
 
   /* Then check if the entire relation is allowed. */
-  for (int i = 0; i < ARRAY_LENGTH(g_pg_catalog_allowed_rels); i++)
+  if (is_member_of(rel_name, g_pg_catalog_allowed_rels, ARRAY_LENGTH(g_pg_catalog_allowed_rels)))
   {
-    if (strcmp(g_pg_catalog_allowed_rels[i], rel_name) == 0)
-    {
-      pfree(rel_name);
-      return true;
-    }
+    pfree(rel_name);
+    return true;
   }
 
   /* Otherwise specific selected columns must be checked against the allow-list. */

src/query/anonymization.c

@@ -591,7 +591,7 @@ static bool collect_seed_material(Node *node, CollectMaterialContext *context)
   if (IsA(node, FuncExpr))
   {
     FuncExpr *func_expr = (FuncExpr *)node;
-    if (!is_allowed_cast(func_expr->funcid))
+    if (!is_allowed_cast(func_expr))
     {
       char *func_name = get_func_name(func_expr->funcid);
       if (func_name)

src/query/validation.c

@@ -271,10 +271,11 @@ static void verify_bucket_expression(Node *node)
   if (IsA(node, FuncExpr))
   {
     FuncExpr *func_expr = (FuncExpr *)node;
-    if (is_allowed_cast(func_expr->funcid))
+    if (is_allowed_cast(func_expr))
     {
       Assert(list_length(func_expr->args) == 1); /* All allowed casts require exactly one argument. */
       verify_bucket_expression(linitial(func_expr->args));
+      return;
     }
 
     if (!is_allowed_function(func_expr->funcid))

test/expected/datetime.out

@@ -97,3 +97,17 @@ SELECT count(*) FROM test_datetime WHERE date_part('century', ts) = 21;
      9
 (1 row)
 
+-- Datetime extract cast to integer
+SELECT cast(extract(minute from ts) as integer) as extract FROM test_datetime GROUP BY 1;
+ extract 
+---------
+       0
+(1 row)
+
+-- Allowed despite the cast being rounding
+SELECT cast(extract(second from ts) as integer) as extract FROM test_datetime GROUP BY 1;
+ extract 
+---------
+       0
+(1 row)
+

test/sql/datetime.sql

@@ -58,3 +58,8 @@ SELECT tz, count(*) FROM test_datetime GROUP BY 1;
 SELECT count(*) FROM test_datetime WHERE date_trunc('year', ts) = '2012-01-01'::timestamp;
 SELECT count(*) FROM test_datetime WHERE extract(century from ts) = 21;
 SELECT count(*) FROM test_datetime WHERE date_part('century', ts) = 21;
+
+-- Datetime extract cast to integer
+SELECT cast(extract(minute from ts) as integer) as extract FROM test_datetime GROUP BY 1;
+-- Allowed despite the cast being rounding
+SELECT cast(extract(second from ts) as integer) as extract FROM test_datetime GROUP BY 1;