Move the global ext_dml_init_hook/ext_dml_fini_hook to AM scoped

The hooks ext_dml_init_hook and ext_dml_fini_hook are used to initialize
and cleanup resources for table DML operation. It's table AM specific.

The current behavior uses global function pointers. The side effect is
that the hooks must check whether the relation is expected and call
the hook chain for other table AM implementations.

This commit moves the hooks to the table AM structure, so the above
assumption could be discarded.

Author: gfphoenix78

contrib/pax_storage/src/cpp/access/pax_access_handle.cc

@@ -412,8 +412,6 @@ void CCPaxAccessMethod::FinishBulkInsert(Relation relation, int options) {
 }
 
 void CCPaxAccessMethod::ExtDmlInit(Relation rel, CmdType operation) {
-  if (!RELATION_IS_PAX(rel)) return;
-
   CBDB_TRY();
   { pax::CPaxDmlStateLocal::Instance()->InitDmlState(rel, operation); }
   CBDB_CATCH_DEFAULT();
@@ -422,8 +420,6 @@ void CCPaxAccessMethod::ExtDmlInit(Relation rel, CmdType operation) {
 }
 
 void CCPaxAccessMethod::ExtDmlFini(Relation rel, CmdType operation) {
-  if (!RELATION_IS_PAX(rel)) return;
-
   CBDB_TRY();
   { pax::CPaxDmlStateLocal::Instance()->FinishDmlState(rel, operation); }
   CBDB_CATCH_DEFAULT();
@@ -791,6 +787,8 @@ static const TableAmRoutine kPaxColumnMethods = {
     .scan_sample_next_block = pax::CCPaxAccessMethod::ScanSampleNextBlock,
     .scan_sample_next_tuple = pax::CCPaxAccessMethod::ScanSampleNextTuple,
 
+    .dml_init = pax::CCPaxAccessMethod::ExtDmlInit,
+    .dml_fini = pax::CCPaxAccessMethod::ExtDmlFini,
     .amoptions = paxc::PaxAccessMethod::AmOptions,
     .swap_relation_files = paxc::PaxAccessMethod::SwapRelationFiles,
     .validate_column_encoding_clauses =
@@ -1191,9 +1189,6 @@ void _PG_init(void) {  // NOLINT
   prev_object_access_hook = object_access_hook;
   object_access_hook = PaxObjectAccessHook;
 
-  ext_dml_init_hook = pax::CCPaxAccessMethod::ExtDmlInit;
-  ext_dml_finish_hook = pax::CCPaxAccessMethod::ExtDmlFini;
-
   prev_ProcessUtilit_hook = ProcessUtility_hook;
   ProcessUtility_hook = paxProcessUtility;
 

contrib/pax_storage/src/cpp/access/pax_access_handle.h

@@ -197,6 +197,3 @@ class CCPaxAccessMethod final {
 };
 
 }  // namespace pax
-
-extern ext_dml_func_hook_type ext_dml_init_hook;
-extern ext_dml_func_hook_type ext_dml_finish_hook;

src/backend/access/aocs/aocsam_handler.c

@@ -270,7 +270,7 @@ remove_dml_state(const Oid relationOid)
  *
  * This function should be called exactly once per relation.
  */
-void
+static void
 aoco_dml_init(Relation relation, CmdType operation)
 {
 	init_aoco_dml_states();
@@ -280,7 +280,7 @@ aoco_dml_init(Relation relation, CmdType operation)
 /*
  * This function should be called exactly once per relation.
  */
-void
+static void
 aoco_dml_finish(Relation relation, CmdType operation)
 {
 	AOCODMLState *state;
@@ -2658,6 +2658,8 @@ static TableAmRoutine ao_column_methods = {
 	.scan_sample_next_block = aoco_scan_sample_next_block,
 	.scan_sample_next_tuple = aoco_scan_sample_next_tuple,
 
+	.dml_init = aoco_dml_init,
+	.dml_fini = aoco_dml_finish,
 	.amoptions = ao_amoptions,
 	.swap_relation_files = aoco_swap_relation_files,
 	.validate_column_encoding_clauses = aoco_validate_column_encoding_clauses,

src/backend/access/appendonly/appendonlyam_handler.c

@@ -228,7 +228,7 @@ remove_dml_state(const Oid relationOid)
  *
  * This function should be called exactly once per relation.
  */
-void
+static void
 appendonly_dml_init(Relation relation, CmdType operation)
 {
 	init_appendonly_dml_states();
@@ -238,7 +238,7 @@ appendonly_dml_init(Relation relation, CmdType operation)
 /*
  * This function should be called exactly once per relation.
  */
-void
+static void
 appendonly_dml_finish(Relation relation, CmdType operation)
 {
 	AppendOnlyDMLState *state;
@@ -2360,6 +2360,8 @@ static const TableAmRoutine ao_row_methods = {
 	.scan_sample_next_block = appendonly_scan_sample_next_block,
 	.scan_sample_next_tuple = appendonly_scan_sample_next_tuple,
 
+	.dml_init = appendonly_dml_init,
+	.dml_fini = appendonly_dml_finish,
 	.amoptions = ao_amoptions,
 	.swap_relation_files = appendonly_swap_relation_files,
 };

src/backend/access/table/tableamapi.c

@@ -104,6 +104,8 @@ GetTableAmRoutine(Oid amhandler)
 		   (routine->scan_bitmap_next_tuple == NULL));
 	Assert(routine->scan_sample_next_block != NULL);
 	Assert(routine->scan_sample_next_tuple != NULL);
+	/* dml_init and dml_fini should both set or ignored */
+	Assert((routine->dml_init != NULL) == (routine->dml_fini != NULL));
 
 	return routine;
 }

src/backend/commands/copyfrom.c

@@ -2007,18 +2007,7 @@ CopyFrom(CopyFromState cstate)
 
 	CopyInitDataParser(cstate);
 
-	/*
-	 * GPDB_12_MERGE_FIXME: We still have to perform the initialization
-	 * here for AO relations. It is preferreable by all means to perform the
-	 * initialization via the table AP API, however it simply does not
-	 * provide a good enough interface for this yet.
-	 */
-	if (RelationIsAoRows(resultRelInfo->ri_RelationDesc))
-		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);
+	table_dml_init(resultRelInfo->ri_RelationDesc, CMD_INSERT);
 
 	for (;;)
 	{
@@ -2451,18 +2440,7 @@ CopyFrom(CopyFromState cstate)
 	if (bistate != NULL)
 		FreeBulkInsertState(bistate);
 
-	/*
-	 * GPDB_12_MERGE_FIXME: We still have to perform the finishment
-	 * here for AO relations. It is preferreable by all means to perform the
-	 * finishment via the table AP API, however it simply does not
-	 * provide a good enough interface for this yet.
-	 */
-	if (RelationIsAoRows(resultRelInfo->ri_RelationDesc))
-		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);
+	table_dml_fini(resultRelInfo->ri_RelationDesc, CMD_INSERT);
 
 	MemoryContextSwitchTo(oldcontext);
 

src/backend/commands/createas.c

@@ -838,21 +838,8 @@ CreateIntoRelDestReceiver(IntoClause *intoClause)
 static void
 intorel_startup_dummy(DestReceiver *self, int operation, TupleDesc typeinfo)
 {
-	/*
-	 * In PostgreSQL, this is a no-op, but in GPDB, AO relations do need some
-	 * initialization of state, because the tableam API does not provide a
-	 * good enough interface for handling with this later, we need to
-	 * specifically all the init at start up.
-	 */
-
 	/* See intorel_initplan() for explanation */
-
-	if (RelationIsAoRows(((DR_intorel *)self)->rel))
-		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);
+	table_dml_init(((DR_intorel *)self)->rel, CMD_INSERT);
 }
 
 /*

src/backend/commands/matview.c

@@ -966,12 +966,7 @@ transientrel_startup(DestReceiver *self, int operation, TupleDesc typeinfo)
 	myState->bistate = GetBulkInsertState();
 	myState->processed = 0;
 
-	if (RelationIsAoRows(myState->transientrel))
-		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);
+	table_dml_init(myState->transientrel, CMD_INSERT);
 
 	/*
 	 * Valid smgr_targblock implies something already wrote to the relation.

src/backend/commands/tablecmds.c

@@ -7560,12 +7560,8 @@ ATRewriteTable(AlteredTableInfo *tab, Oid OIDNewHeap, LOCKMODE lockmode)
 		snapshot = RegisterSnapshot(GetLatestSnapshot());
 		scan = table_beginscan(oldrel, snapshot, 0, NULL);
 
-		if (newrel && RelationIsAoRows(newrel))
-			appendonly_dml_init(newrel, CMD_INSERT);
-		else if (newrel && RelationIsAoCols(newrel))
-			aoco_dml_init(newrel, CMD_INSERT);
-		else if (newrel && ext_dml_init_hook)
-			ext_dml_init_hook(newrel, CMD_INSERT);
+		if (newrel)
+			table_dml_init(newrel, CMD_INSERT);
 
 		/*
 		 * Switch to per-tuple memory context and reset it for each tuple

src/backend/executor/execPartition.c

@@ -916,12 +916,7 @@ ExecInitPartitionInfo(ModifyTableState *mtstate, EState *estate,
 		lappend(estate->es_tuple_routing_result_relations,
 				leaf_part_rri);
 
-	if (RelationIsAoRows(leaf_part_rri->ri_RelationDesc))
-		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);
+	table_dml_init(leaf_part_rri->ri_RelationDesc, mtstate->operation);
 
 	MemoryContextSwitchTo(oldcxt);
 
@@ -1234,12 +1229,7 @@ ExecCleanupTupleRouting(ModifyTableState *mtstate,
 		 * Only leaf node can have a valid access method.  If we find an
 		 * appendoptimized table, ensure the DML operation is finished.
 		 */
-		if (RelationIsAoRows(resultRelInfo->ri_RelationDesc))
-			appendonly_dml_finish(resultRelInfo->ri_RelationDesc, mtstate->operation);
-		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);
+		table_dml_fini(resultRelInfo->ri_RelationDesc, mtstate->operation);
 
 		ExecCloseIndices(resultRelInfo);
 		table_close(resultRelInfo->ri_RelationDesc, NoLock);