on partition created callback

zilder · zilder · commit 0a3b7e597fd0 · 2016-09-16T17:37:27.000+03:00
diff --git a/expected/pg_pathman.out b/expected/pg_pathman.out
@@ -1768,3 +1768,59 @@ NOTICE:  100 rows copied from test_fkey_0
               10
 (1 row)
 
+/* Check callbacks */
+CREATE TABLE log(id serial, message text);
+CREATE FUNCTION abc_on_partition_created_callback(args jsonb)
+RETURNS VOID AS $$
+DECLARE
+  start_value TEXT := args->>'start';
+  end_value   TEXT := args::jsonb->'end';
+BEGIN
+  -- raise notice 'callback: %', args->start;
+  INSERT INTO log(message)
+  VALUES (start_value || '-' || end_value);
+END
+$$ language plpgsql;
+CREATE TABLE abc(a serial, b int);
+SELECT create_range_partitions('abc', 'a', 1, 100, 2);
+NOTICE:  sequence "abc_seq" does not exist, skipping
+ create_range_partitions 
+-------------------------
+                       2
+(1 row)
+
+SELECT set_callback('abc', 'abc_on_partition_created_callback');
+ set_callback 
+--------------
+ 
+(1 row)
+
+INSERT INTO abc VALUES (123, 1);
+INSERT INTO abc VALUES (223, 1);
+SELECT append_range_partition('abc');
+ append_range_partition 
+------------------------
+ public.abc_4
+(1 row)
+
+SELECT prepend_range_partition('abc');
+ prepend_range_partition 
+-------------------------
+ public.abc_5
+(1 row)
+
+SELECT add_range_partition('abc', 401, 501);
+ add_range_partition 
+---------------------
+ public.abc_6
+(1 row)
+
+SELECT message FROM log;
+ message 
+---------
+ 201-301
+ 301-401
+ -99-1
+ 401-501
+(4 rows)
+
diff --git a/init.sql b/init.sql
@@ -31,23 +31,17 @@ CREATE TABLE IF NOT EXISTS @extschema@.pathman_config (
  *		partrel - regclass (relation type, stored as Oid)
  *		enable_parent - add parent table to plan
  *		auto - enable automatic partition creation
+ *		callback - 
  */
 CREATE TABLE IF NOT EXISTS @extschema@.pathman_config_params (
 	partrel			REGCLASS NOT NULL PRIMARY KEY,
 	enable_parent	BOOLEAN NOT NULL DEFAULT TRUE,
-	auto			BOOLEAN NOT NULL DEFAULT TRUE
+	auto			BOOLEAN NOT NULL DEFAULT TRUE,
+	callback		REGPROCEDURE
 );
 CREATE UNIQUE INDEX i_pathman_config_params
 ON @extschema@.pathman_config_params(partrel);
 
-CREATE TYPE partition AS (
-	parent			REGCLASS,
-	parttype		INTEGER,
-	child			REGCLASS,
-	start_value		TEXT,
-	end_value		TEXT
-);
-
 /*
  * Invalidate relcache every time someone changes parameters config.
  */
@@ -101,7 +95,7 @@ LANGUAGE plpgsql;
 CREATE OR REPLACE FUNCTION @extschema@.pathman_set_param(
 	relation	REGCLASS,
 	param		TEXT,
-	value		BOOLEAN)
+	value		ANYELEMENT)
 RETURNS VOID AS
 $$
 BEGIN
@@ -161,6 +155,19 @@ END
 $$
 LANGUAGE plpgsql;
 
+/*
+ * Set partition creation callback
+ */
+CREATE OR REPLACE FUNCTION @extschema@.set_callback(relation REGCLASS, callback REGPROC)
+RETURNS VOID AS
+$$
+BEGIN
+	PERFORM @extschema@.validate_on_partition_created_callback(callback);
+	PERFORM @extschema@.pathman_set_param(relation, 'callback', callback);
+END
+$$
+LANGUAGE plpgsql;
+
 /*
  * Show all existing concurrent partitioning tasks.
  */
@@ -710,3 +717,23 @@ LANGUAGE C STRICT;
 CREATE OR REPLACE FUNCTION @extschema@.get_rel_tablespace_name(relation REGCLASS)
 RETURNS TEXT AS 'pg_pathman', 'get_rel_tablespace_name'
 LANGUAGE C STRICT;
+
+/*
+ * Checks that callback function meets specific requirements. Particularly it
+ * must have the only JSONB argument and VOID return type
+ */
+CREATE OR REPLACE FUNCTION @extschema@.validate_on_partition_created_callback(callback REGPROC)
+RETURNS VOID AS 'pg_pathman', 'validate_on_partition_created_callback'
+LANGUAGE C STRICT;
+
+/*
+ * Builds JSONB object containing new partition parameters and invoke the
+ * callback
+ */
+CREATE OR REPLACE FUNCTION @extschema@.invoke_on_partition_created_callback(
+	parent		REGCLASS,
+	partition	REGCLASS,
+	start_value	ANYELEMENT,
+	end_value	ANYELEMENT)
+RETURNS VOID AS 'pg_pathman', 'invoke_on_partition_created_callback'
+LANGUAGE C STRICT;
diff --git a/range.sql b/range.sql
@@ -504,7 +504,6 @@ BEGIN
 	if NOT tablespace IS NULL THEN
 		v_create_table_query := v_create_table_query || ' TABLESPACE ' ||tablespace;
 	END IF;
-	RAISE NOTICE 'query: %', v_create_table_query;
 
 	EXECUTE format(v_create_table_query,
 				   v_child_relname,
@@ -519,6 +518,10 @@ BEGIN
 													 p_end_value));
 
 	PERFORM @extschema@.copy_foreign_keys(parent_relid, v_child_relname::REGCLASS);
+	PERFORM @extschema@.invoke_on_partition_created_callback(parent_relid,
+															 v_child_relname,
+															 p_start_value,
+															 p_end_value);
 
 	RETURN v_child_relname;
 END
@@ -1234,7 +1237,6 @@ BEGIN
 END
 $$ LANGUAGE plpgsql;
 
-
 /*
  * Construct CHECK constraint condition for a range partition.
  */
diff --git a/sql/pg_pathman.sql b/sql/pg_pathman.sql
@@ -665,3 +665,27 @@ SELECT create_hash_partitions('test_fkey', 'id', 10);
 INSERT INTO test_fkey VALUES(1, 'wrong');
 INSERT INTO test_fkey VALUES(1, 'test');
 SELECT drop_partitions('test_fkey');
+
+/* Check callbacks */
+CREATE TABLE log(id serial, message text);
+
+CREATE OR REPLACE FUNCTION abc_on_partition_created_callback(args jsonb)
+RETURNS VOID AS $$
+DECLARE
+	start_value TEXT := args->>'start';
+	end_value   TEXT := args::jsonb->'end';
+BEGIN
+	INSERT INTO log(message)
+	VALUES (start_value || '-' || end_value);
+END
+$$ language plpgsql;
+
+CREATE TABLE abc(a serial, b int);
+SELECT create_range_partitions('abc', 'a', 1, 100, 2);
+SELECT set_callback('abc', 'abc_on_partition_created_callback');
+INSERT INTO abc VALUES (123, 1);
+INSERT INTO abc VALUES (223, 1);
+SELECT append_range_partition('abc');
+SELECT prepend_range_partition('abc');
+SELECT add_range_partition('abc', 401, 501);
+SELECT message FROM log;
diff --git a/src/pathman.h b/src/pathman.h
@@ -55,10 +55,11 @@
  * Definitions for the "pathman_config_params" table
  */
 #define PATHMAN_CONFIG_PARAMS						"pathman_config_params"
-#define Natts_pathman_config_params					3
+#define Natts_pathman_config_params					4
 #define Anum_pathman_config_params_partrel			1	/* primary key */
 #define Anum_pathman_config_params_enable_parent	2	/* include parent into plan */
 #define Anum_pathman_config_params_auto				3	/* auto partitions creation */
+#define Anum_pathman_config_params_callback			4	/* auto partitions creation */
 
 /*
  * Cache current PATHMAN_CONFIG relid (set during load_config()).
diff --git a/src/pl_funcs.c b/src/pl_funcs.c
@@ -13,11 +13,14 @@
 #include "relation_info.h"
 #include "utils.h"
 #include "xact_handling.h"
+#include "fmgr.h"
 
 #include "access/htup_details.h"
 #include "access/nbtree.h"
 #include "access/xact.h"
 #include "catalog/indexing.h"
+#include "catalog/pg_type.h"
+#include "catalog/pg_proc.h"
 #include "commands/sequence.h"
 #include "commands/tablespace.h"
 #include "miscadmin.h"
@@ -28,6 +31,8 @@
 #include "utils/lsyscache.h"
 #include "utils/syscache.h"
 #include "utils/typcache.h"
+#include "utils/jsonb.h"
+#include "utils/fmgroids.h"
 
 
 /* declarations */
@@ -57,7 +62,8 @@ PG_FUNCTION_INFO_V1( lock_partitioned_relation );
 PG_FUNCTION_INFO_V1( prevent_relation_modification );
 PG_FUNCTION_INFO_V1( debug_capture );
 PG_FUNCTION_INFO_V1( get_rel_tablespace_name );
-
+PG_FUNCTION_INFO_V1( validate_on_partition_created_callback );
+PG_FUNCTION_INFO_V1( invoke_on_partition_created_callback );
 
 static void on_partitions_created_internal(Oid partitioned_table, bool add_callbacks);
 static void on_partitions_updated_internal(Oid partitioned_table, bool add_callbacks);
@@ -776,3 +782,87 @@ get_rel_tablespace_name(PG_FUNCTION_ARGS)
 	result = get_tablespace_name(tablespace_id);
 	PG_RETURN_TEXT_P(cstring_to_text(result));
 }
+
+/*
+ * Checks that callback function meets specific requirements. Particularly it
+ * must have the only JSONB argument and VOID return type
+ */
+Datum
+validate_on_partition_created_callback(PG_FUNCTION_ARGS)
+{
+	HeapTuple	tp;
+	Oid			callback = PG_GETARG_OID(0);
+	Form_pg_proc functup;
+
+	tp = SearchSysCache1(PROCOID, ObjectIdGetDatum(callback));
+	if (!HeapTupleIsValid(tp))
+		elog(ERROR, "cache lookup failed for function %u", callback);
+	functup = (Form_pg_proc) GETSTRUCT(tp);
+
+	if (functup->pronargs != 1 || functup->proargtypes.values[0] != JSONBOID ||
+		functup->prorettype != VOIDOID)
+		elog(ERROR,
+			 "Callback function must have only one JSNOB argument "
+			 "and return VOID");
+
+	ReleaseSysCache(tp);
+	PG_RETURN_VOID();
+}
+
+/*
+ * Builds JSONB object containing new partition parameters and invoke the
+ * callback
+ */
+Datum
+invoke_on_partition_created_callback(PG_FUNCTION_ARGS)
+{
+	char   *json;
+	Datum	jsonb;
+	Oid		parent_oid = PG_GETARG_OID(0);
+	Oid		partition_oid = PG_GETARG_OID(1);
+	Oid		type = get_fn_expr_argtype(fcinfo->flinfo, 2);
+	Datum	start_value = PG_GETARG_DATUM(2);
+	Datum	end_value = PG_GETARG_DATUM(3);
+	const PartRelationInfo   *prel;
+
+	if ((prel = get_pathman_relation_info(parent_oid)) == NULL)
+		elog(ERROR,
+			 "Relation %s isn't partitioned by pg_pathman",
+			 get_rel_name(parent_oid));
+
+	/* If there is no callback function specified then we're done */
+	if (!prel->callback)
+		PG_RETURN_VOID();
+	
+	/* Convert ANYELEMENT arguments to jsonb */
+	start_value = convert_to_jsonb(start_value, type);
+	end_value = convert_to_jsonb(end_value, type);
+
+	/*
+	 * Build jsonb object to pass into callback
+	 *
+	 * XXX it would be nice to have this rewrited with pushJsonbValue() to get
+	 * rid of string formatting and parsing. See jsonb_build_object() for
+	 * example
+	 */
+	json = psprintf("{"
+		"\"parent\": %u,"
+		"\"partition\": %u,"
+		"\"part_type\": %u,"
+		"\"start\": %s,"
+		"\"end\": %s,"
+		"\"value_type\": %u}",
+		parent_oid,
+		partition_oid,
+		prel->parttype,
+		datum_to_cstring(start_value, JSONBOID),
+		datum_to_cstring(end_value, JSONBOID),
+		type
+	);
+	jsonb = OidFunctionCall1(F_JSONB_IN, CStringGetDatum(json));
+
+	/* Invoke callback */
+	OidFunctionCall1(prel->callback, JsonbGetDatum(jsonb));
+
+	PG_RETURN_JSONB(jsonb);
+}
diff --git a/src/relation_info.c b/src/relation_info.c
@@ -165,12 +165,14 @@ refresh_pathman_relation_info(Oid relid,
 	{
 		prel->enable_parent = param_values[Anum_pathman_config_params_enable_parent - 1];
 		prel->auto_partition = param_values[Anum_pathman_config_params_auto - 1];
+		prel->callback = param_values[Anum_pathman_config_params_callback - 1];
 	}
 	/* Else set default values if they cannot be found */
 	else
 	{
 		prel->enable_parent = false;
 		prel->auto_partition = true;
+		prel->callback = InvalidOid;
 	}
 
 	/* We've successfully built a cache entry */
diff --git a/src/relation_info.h b/src/relation_info.h
@@ -49,6 +49,7 @@ typedef struct
 	bool			valid;			/* is this entry valid? */
 	bool			enable_parent;	/* include parent to the plan */
 	bool			auto_partition; /* auto partition creation */
+	Oid				callback;		/* callback for partition creation */
 
 	uint32			children_count;
 	Oid			   *children;		/* Oids of child partitions */
diff --git a/src/utils.c b/src/utils.c
diff --git a/src/utils.h b/src/utils.h

Original file line number	Diff line number	Diff line change
`@@ -165,12 +165,14 @@ refresh_pathman_relation_info(Oid relid,`
`165`	`165`	`{`
`166`	`166`	`prel->enable_parent = param_values[Anum_pathman_config_params_enable_parent - 1];`
`167`	`167`	`prel->auto_partition = param_values[Anum_pathman_config_params_auto - 1];`
	`168`	`+ prel->callback = param_values[Anum_pathman_config_params_callback - 1];`
`168`	`169`	`}`
`169`	`170`	`/* Else set default values if they cannot be found */`
`170`	`171`	`else`
`171`	`172`	`{`
`172`	`173`	`prel->enable_parent = false;`
`173`	`174`	`prel->auto_partition = true;`
	`175`	`+ prel->callback = InvalidOid;`
`174`	`176`	`}`
`175`	`177`
`176`	`178`	`/* We've successfully built a cache entry */`