Support extract(... from ...)

src/query/allowed_objects.c

@@ -1,6 +1,7 @@
 #include "postgres.h"
 
 #include "access/sysattr.h"
+#include "utils/fmgroids.h"
 #include "utils/fmgrtab.h"
 #include "utils/lsyscache.h"
 #include "utils/memutils.h"
@@ -57,6 +58,12 @@ static const FunctionByName g_allowed_builtins[] = {
     /* date_trunc */
     (FunctionByName){.name = "timestamptz_trunc", .primary_arg = 1},
     (FunctionByName){.name = "timestamp_trunc", .primary_arg = 1},
+    /* extract & date_part*/
+    (FunctionByName){.name = "extract_date", .primary_arg = 1},
+    (FunctionByName){.name = "extract_timestamp", .primary_arg = 1},
+    (FunctionByName){.name = "extract_timestamptz", .primary_arg = 1},
+    (FunctionByName){.name = "timestamp_part", .primary_arg = 1},
+    (FunctionByName){.name = "timestamptz_part", .primary_arg = 1},
     /**/
 };
 
@@ -73,7 +80,19 @@ static const char *const g_implicit_range_builtins_untrusted[] = {
 
 /* Some allowed functions don't appear in the builtins catalog, so we must allow them manually by OID. */
 #define F_NUMERIC_ROUND_INT 1708
-static const FunctionByOid g_allowed_builtins_extra[] = {(FunctionByOid){.funcid = F_NUMERIC_ROUND_INT, .primary_arg = 0}};
+/*
+ * `date_part` for `date` is a SQL builtin and doesn't show up in `fmgr_isbuiltin`.
+ *  PG 14 has the define, but PG 13 doesn't.
+ */
+#if PG_MAJORVERSION_NUM < 14
+#define F_DATE_PART_TEXT_DATE 1384
+#endif
+
+static const FunctionByOid g_allowed_builtins_extra[] = {
+    (FunctionByOid){.funcid = F_NUMERIC_ROUND_INT, .primary_arg = 0},
+    (FunctionByOid){.funcid = F_DATE_PART_TEXT_DATE, .primary_arg = 1},
+    /**/
+};
 
 typedef struct AllowedCols
 {

src/query/anonymization.c

@@ -574,6 +574,15 @@ typedef struct CollectMaterialContext
   char material[MAX_SEED_MATERIAL_SIZE];
 } CollectMaterialContext;
 
+static void normalize_function_name(char *func_name)
+{
+  if (strcmp(func_name, "date_part") == 0)
+  {
+    // Not reallocing the `func_name`, because the normalized string is shorter.
+    strcpy(func_name, "extract");
+  }
+}
+
 static bool collect_seed_material(Node *node, CollectMaterialContext *context)
 {
   if (node == NULL)
@@ -587,7 +596,7 @@ static bool collect_seed_material(Node *node, CollectMaterialContext *context)
       char *func_name = get_func_name(func_expr->funcid);
       if (func_name)
       {
-        /* TODO: Normalize function names. */
+        normalize_function_name(func_name);
         append_seed_material(context->material, func_name, ',');
         pfree(func_name);
       }

test/expected/datetime.out

@@ -85,3 +85,15 @@ SELECT count(*) FROM test_datetime WHERE date_trunc('year', ts) = '2012-01-01'::
     11
 (1 row)
 
+SELECT count(*) FROM test_datetime WHERE extract(century from ts) = 21;
+ count 
+-------
+     9
+(1 row)
+
+SELECT count(*) FROM test_datetime WHERE date_part('century', ts) = 21;
+ count 
+-------
+     9
+(1 row)
+

test/expected/validation.out

@@ -126,14 +126,21 @@ GROUP BY 1, 2, 3, 4;
 -----------+-----------+-----------+-----------
 (0 rows)
 
+-- `as extract` ensures that the column is aliased consistently in PG 13 and 14.
 SELECT
   date_trunc('year', last_seen),
   date_trunc('year', last_seen_tz),
-  date_trunc('year', birthday)
+  date_trunc('year', birthday),
+  extract(month from last_seen) as extract,
+  extract(month from last_seen_tz) as extract,
+  extract(month from birthday) as extract,
+  date_part('month', last_seen) as date_part,
+  date_part('month', last_seen_tz) as date_part,
+  date_part('month', birthday) as date_part
 FROM test_validation
-GROUP BY 1, 2, 3;
- date_trunc | date_trunc | date_trunc 
-------------+------------+------------
+GROUP BY 1, 2, 3, 4, 5, 6, 7, 8, 9;
+ date_trunc | date_trunc | date_trunc | extract | extract | extract | date_part | date_part | date_part 
+------------+------------+------------+---------+---------+---------+-----------+-----------+-----------
 (0 rows)
 
 -- Allow all functions post-anonymization.
@@ -452,6 +459,10 @@ SELECT date_trunc('year', lunchtime) FROM test_validation GROUP BY 1;
 ERROR:  [PG_DIFFIX] Unsupported function used for generalization.
 LINE 1: SELECT date_trunc('year', lunchtime) FROM test_validation GR...
                ^
+SELECT extract(hour from lunchtime) FROM test_validation GROUP BY 1;
+ERROR:  [PG_DIFFIX] Unsupported function used for generalization.
+LINE 1: SELECT extract(hour from lunchtime) FROM test_validation GRO...
+               ^
 -- Get rejected because of averaging opportunity
 SELECT  date_trunc('year', last_seen_tz, 'EST') FROM test_validation GROUP BY 1;
 ERROR:  [PG_DIFFIX] Unsupported function used for generalization.

test/sql/datetime.sql

@@ -56,3 +56,5 @@ SELECT tz, count(*) FROM test_datetime GROUP BY 1;
 
 -- Datetime filtering
 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;

test/sql/validation.sql

@@ -93,12 +93,19 @@ SELECT
 FROM test_validation
 GROUP BY 1, 2, 3, 4;
 
+-- `as extract` ensures that the column is aliased consistently in PG 13 and 14.
 SELECT
   date_trunc('year', last_seen),
   date_trunc('year', last_seen_tz),
-  date_trunc('year', birthday)
+  date_trunc('year', birthday),
+  extract(month from last_seen) as extract,
+  extract(month from last_seen_tz) as extract,
+  extract(month from birthday) as extract,
+  date_part('month', last_seen) as date_part,
+  date_part('month', last_seen_tz) as date_part,
+  date_part('month', birthday) as date_part
 FROM test_validation
-GROUP BY 1, 2, 3;
+GROUP BY 1, 2, 3, 4, 5, 6, 7, 8, 9;
 
 -- Allow all functions post-anonymization.
 SELECT 2 * length(city) FROM test_validation GROUP BY city;
@@ -224,6 +231,7 @@ SELECT COUNT(*) FROM test_validation GROUP BY substr('aaaa', 1, 2);
 
 -- Get rejected because of lack of interval support
 SELECT date_trunc('year', lunchtime) FROM test_validation GROUP BY 1;
+SELECT extract(hour from lunchtime) FROM test_validation GROUP BY 1;
 
 -- Get rejected because of averaging opportunity
 SELECT  date_trunc('year', last_seen_tz, 'EST') FROM test_validation GROUP BY 1;