From 53e55ba44cdab0be6165c9741943b3a79d63f171 Mon Sep 17 00:00:00 2001 From: Hao Wu Date: Thu, 20 Apr 2023 22:19:52 +0800 Subject: [PATCH] Add dml hook for extensions Add dml hooks to initialize or finish modifying relations, like aoco_dml_init/aoco_dml_finish, for extensions. --- src/backend/commands/copyfrom.c | 4 ++++ src/backend/commands/createas.c | 2 ++ src/backend/commands/matview.c | 2 ++ src/backend/commands/tablecmds.c | 2 ++ src/backend/executor/execPartition.c | 6 +++++- src/backend/executor/nodeModifyTable.c | 6 ++++++ src/include/access/tableam.h | 9 +++++++++ 7 files changed, 30 insertions(+), 1 deletion(-) diff --git a/src/backend/commands/copyfrom.c b/src/backend/commands/copyfrom.c index ba05e78b3d7..acb01d91ec6 100644 --- 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 index 8fec01684c0..aa4c7f53155 100644 --- 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 index 5f292bc8aec..a662bc1b0e4 100644 --- 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 index eac95eeb920..af03e95ff29 100644 --- 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 index b6e10662a89..fc2b4adb1f7 100644 --- 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 index 7d98b69d345..8bac299db77 100644 --- 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 index 61c13537c3c..aabec3dc474 100644 --- a/src/include/access/tableam.h +++ b/src/include/access/tableam.h @@ -2193,4 +2193,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 */