Add dml hook for extensions

Add dml hooks to initialize or finish modifying relations, like aoco_dml_init/aoco_dml_finish, for extensions.
apache · Dec 12, 2023 · 92496d8 · 92496d8
1 parent a6a1992
commit 92496d8
Show file tree

Hide file tree

Showing 7 changed files with 30 additions and 1 deletion.
diff --git a/src/backend/commands/copyfrom.c b/src/backend/commands/copyfrom.c
@@ -1093,6 +1093,8 @@ CopyFrom(CopyFromState cstate)
 		appendonly_dml_init(resultRelInfo->ri_RelationDesc, CMD_INSERT);
 	else if (RelationIsAoCols(resultRelInfo->ri_RelationDesc))
 		aoco_dml_init(resultRelInfo->ri_RelationDesc, CMD_INSERT);
+	else if (ext_dml_init_hook)
+		ext_dml_init_hook(resultRelInfo->ri_RelationDesc, CMD_INSERT);
 
 	for (;;)
 	{
@@ -1533,6 +1535,8 @@ CopyFrom(CopyFromState cstate)
 		appendonly_dml_finish(resultRelInfo->ri_RelationDesc, CMD_INSERT);
 	else if (RelationIsAoCols(resultRelInfo->ri_RelationDesc))
 		aoco_dml_finish(resultRelInfo->ri_RelationDesc, CMD_INSERT);
+	else if (ext_dml_finish_hook)
+		ext_dml_finish_hook(resultRelInfo->ri_RelationDesc, CMD_INSERT);
 
 	MemoryContextSwitchTo(oldcontext);
 

diff --git a/src/backend/commands/createas.c b/src/backend/commands/createas.c
@@ -804,6 +804,8 @@ intorel_startup_dummy(DestReceiver *self, int operation, TupleDesc typeinfo)
 		appendonly_dml_init(((DR_intorel *)self)->rel, CMD_INSERT);
 	else if (RelationIsAoCols(((DR_intorel *)self)->rel))
 		aoco_dml_init(((DR_intorel *)self)->rel, CMD_INSERT);
+	else if (ext_dml_init_hook)
+		ext_dml_init_hook(((DR_intorel *)self)->rel, CMD_INSERT);
 }
 
 /*

diff --git a/src/backend/commands/matview.c b/src/backend/commands/matview.c
@@ -864,6 +864,8 @@ transientrel_startup(DestReceiver *self, int operation, TupleDesc typeinfo)
 		appendonly_dml_init(myState->transientrel, CMD_INSERT);
 	else if (RelationIsAoCols(myState->transientrel))
 		aoco_dml_init(myState->transientrel, CMD_INSERT);
+	else if (ext_dml_init_hook)
+		ext_dml_init_hook(myState->transientrel, CMD_INSERT);
 
 	/*
 	 * Valid smgr_targblock implies something already wrote to the relation.

diff --git a/src/backend/commands/tablecmds.c b/src/backend/commands/tablecmds.c
@@ -7132,6 +7132,8 @@ ATRewriteTable(AlteredTableInfo *tab, Oid OIDNewHeap, LOCKMODE lockmode)
 			Assert(RelationIsAoCols(oldrel));
 			aoco_dml_init(newrel, CMD_INSERT);
 		}
+		else if (newrel && ext_dml_init_hook)
+			ext_dml_init_hook(newrel, CMD_INSERT);
 
 		/*
 		 * Switch to per-tuple memory context and reset it for each tuple

diff --git a/src/backend/executor/execPartition.c b/src/backend/executor/execPartition.c
@@ -864,6 +864,8 @@ ExecInitPartitionInfo(ModifyTableState *mtstate, EState *estate,
 		appendonly_dml_init(leaf_part_rri->ri_RelationDesc, mtstate->operation);
 	else if (RelationIsAoCols(leaf_part_rri->ri_RelationDesc))
 		aoco_dml_init(leaf_part_rri->ri_RelationDesc, mtstate->operation);
+	else if (ext_dml_init_hook)
+		ext_dml_init_hook(leaf_part_rri->ri_RelationDesc, mtstate->operation);
 
 	MemoryContextSwitchTo(oldcxt);
 
@@ -1178,8 +1180,10 @@ ExecCleanupTupleRouting(ModifyTableState *mtstate,
 		 */
 		if (RelationIsAoRows(resultRelInfo->ri_RelationDesc))
 			appendonly_dml_finish(resultRelInfo->ri_RelationDesc, mtstate->operation);
-		if (RelationIsAoCols(resultRelInfo->ri_RelationDesc))
+		else if (RelationIsAoCols(resultRelInfo->ri_RelationDesc))
 			aoco_dml_finish(resultRelInfo->ri_RelationDesc, mtstate->operation);
+		else if (ext_dml_finish_hook)
+			ext_dml_finish_hook(resultRelInfo->ri_RelationDesc, mtstate->operation);
 
 		ExecCloseIndices(resultRelInfo);
 		table_close(resultRelInfo->ri_RelationDesc, NoLock);

diff --git a/src/backend/executor/nodeModifyTable.c b/src/backend/executor/nodeModifyTable.c
@@ -63,6 +63,8 @@
 #include "utils/lsyscache.h"
 #include "utils/snapmgr.h"
 
+ext_dml_func_hook_type ext_dml_init_hook   = NULL;
+ext_dml_func_hook_type ext_dml_finish_hook = NULL;
 
 typedef struct MTTargetRelLookup
 {
@@ -3069,6 +3071,8 @@ ExecInitModifyTable(ModifyTable *node, EState *estate, int eflags)
 			appendonly_dml_init(resultRelInfo->ri_RelationDesc, operation);
 		else if (RelationIsAoCols(resultRelInfo->ri_RelationDesc))
 			aoco_dml_init(resultRelInfo->ri_RelationDesc, operation);
+		else if (ext_dml_init_hook)
+			ext_dml_init_hook(resultRelInfo->ri_RelationDesc, operation);
 
 		resultRelInfo++;
 		i++;
@@ -3527,6 +3531,8 @@ ExecEndModifyTable(ModifyTableState *node)
 								  node->operation);
 		else if (RelationIsAoCols(resultRelInfo->ri_RelationDesc))
 			aoco_dml_finish(resultRelInfo->ri_RelationDesc, node->operation);
+		else if (ext_dml_finish_hook)
+			ext_dml_finish_hook(resultRelInfo->ri_RelationDesc, node->operation);
 
 		/*
 		 * Cleanup the initialized batch slots. This only matters for FDWs

diff --git a/src/include/access/tableam.h b/src/include/access/tableam.h
@@ -2183,4 +2183,13 @@ extern const TableAmRoutine *GetHeapamTableAmRoutine(void);
 extern bool check_default_table_access_method(char **newval, void **extra,
 											  GucSource source);
 
+/* ----------------------------------------------------------------------------
+ * Hook function to run init/fini for storage extensions
+ * ----------------------------------------------------------------------------
+ */
+enum CmdType;
+typedef void (*ext_dml_func_hook_type) (Relation relation, enum CmdType operation);
+extern PGDLLIMPORT ext_dml_func_hook_type ext_dml_init_hook;
+extern PGDLLIMPORT ext_dml_func_hook_type ext_dml_finish_hook;
+
 #endif							/* TABLEAM_H */