Add new wrapper function for common Table Variable and Temp Table checks

Signed-off-by: Tim Chang <[email protected]>
babelfish-for-postgresql · Sep 26, 2024 · f7796cd · f7796cd
1 parent 7469009
commit f7796cd
Show file tree

Hide file tree

Showing 10 changed files with 40 additions and 21 deletions.
diff --git a/src/backend/access/heap/heapam_visibility.c b/src/backend/access/heap/heapam_visibility.c
@@ -80,6 +80,7 @@
 #include "utils/builtins.h"
 #include "utils/combocid.h"
 #include "utils/snapmgr.h"
+#include "utils/queryenvironment.h"
 
 table_variable_satisfies_visibility_hook_type table_variable_satisfies_visibility_hook = NULL;
 table_variable_satisfies_update_hook_type table_variable_satisfies_update_hook = NULL;
@@ -468,7 +469,7 @@ HeapTupleSatisfiesUpdate(Relation relation, HeapTuple htup, CommandId curcid,
 	HeapTupleHeader tuple = htup->t_data;
 
 	/* See HeapTupleSatisfiesVisibility why */
-	if (sql_dialect == SQL_DIALECT_TSQL && RelationIsBBFTableVariable(relation))
+	if (IsTsqlTableVariable(relation))
 		return table_variable_satisfies_update_hook(htup, curcid, buffer);
 
 	Assert(ItemPointerIsValid(&htup->t_self));
@@ -1177,7 +1178,7 @@ HeapTupleSatisfiesVacuum(Relation relation, HeapTuple htup, TransactionId Oldest
 	HTSV_Result res;
 
 	/* See HeapTupleSatisfiesVisibility why */
-	if (sql_dialect == SQL_DIALECT_TSQL && RelationIsBBFTableVariable(relation))
+	if (IsTsqlTableVariable(relation))
 		return table_variable_satisfies_vacuum_hook(htup, OldestXmin, buffer);
 
 	res = HeapTupleSatisfiesVacuumHorizon(NULL, htup, buffer, &dead_after);
@@ -1219,7 +1220,7 @@ HeapTupleSatisfiesVacuumHorizon(Relation relation, HeapTuple htup, Buffer buffer
 	*dead_after = InvalidTransactionId;
 
 	/* See HeapTupleSatisfiesVisibility why */
-	if (sql_dialect == SQL_DIALECT_TSQL && relation && RelationIsBBFTableVariable(relation))
+	if (IsTsqlTableVariable(relation))
 		return table_variable_satisfies_vacuum_horizon_hook(htup, buffer, dead_after);
 
 	/*
@@ -1789,7 +1790,7 @@ HeapTupleSatisfiesVisibility(Relation relation, HeapTuple htup, Snapshot snapsho
 	* Babelfish extension has a different type of heap table but non-transactional.
 	* It is not worth introducing a new table AM because the new table still uses heap format.
 	*/
-	if (sql_dialect == SQL_DIALECT_TSQL && relation && RelationIsBBFTableVariable(relation))
+	if (IsTsqlTableVariable(relation))
 		return table_variable_satisfies_visibility_hook(htup, snapshot, buffer);
 
 	switch (snapshot->snapshot_type)

diff --git a/src/backend/catalog/catalog.c b/src/backend/catalog/catalog.c
@@ -583,7 +583,7 @@ GetNewRelFileNumber(Oid reltablespace, Relation pg_class, char relpersistence)
 	 */
 	rlocator.backend = backend;
 
-	use_bbf_oid_buffer = (relpersistence == RELPERSISTENCE_TEMP && useTempOidBuffer());
+	use_bbf_oid_buffer = (relpersistence == RELPERSISTENCE_TEMP && UseTempOidBuffer());
 
 	do
 	{

diff --git a/src/backend/catalog/heap.c b/src/backend/catalog/heap.c
@@ -1185,7 +1185,7 @@ heap_create_with_catalog(const char *relname,
 	MultiXactId relminmxid;
 	bool		is_enr = false;
 
-	if (relpersistence == RELPERSISTENCE_TEMP && sql_dialect == SQL_DIALECT_TSQL)
+	if (IsTsqlTempTable(relpersistence))
 	{
 		/*
 		 * in TSQL, temporary table name should start with '#'.
@@ -1413,7 +1413,7 @@ heap_create_with_catalog(const char *relname,
 		 * 
 		 * For temp tables, we use temp OID assignment code here as well.
 		 */
-		if (is_enr && useTempOidBuffer())
+		if (is_enr && UseTempOidBuffer())
 		{
 			Relation	pg_type;
 

diff --git a/src/backend/catalog/pg_constraint.c b/src/backend/catalog/pg_constraint.c
@@ -176,7 +176,7 @@ CreateConstraintEntry(const char *constraintName,
 		values[i] = (Datum) NULL;
 	}
 
-	if (useTempOidBufferForOid(relId) && isTempNamespace(constraintNamespace))
+	if (UseTempOidBufferForOid(relId) && isTempNamespace(constraintNamespace))
 		conOid = GetNewTempOidWithIndex_hook(conDesc, ConstraintOidIndexId,
 									Anum_pg_constraint_oid);
 	else

diff --git a/src/backend/catalog/storage.c b/src/backend/catalog/storage.c
@@ -33,6 +33,7 @@
 #include "storage/smgr.h"
 #include "utils/hsearch.h"
 #include "utils/memutils.h"
+#include "utils/queryenvironment.h"
 #include "utils/rel.h"
 
 /* GUC variables */
@@ -223,7 +224,7 @@ RelationDropStorage(Relation rel)
 	 * However, we need to unlink the files on explicit DROP TABLE command regardless
 	 * if the transaction state is committing or aborting.
 	 */
-	if (sql_dialect == SQL_DIALECT_TSQL && RelationIsBBFTableVariable(rel))
+	if (IsTsqlTableVariable(rel))
 	{
 		pending = (PendingRelDelete *)
 		MemoryContextAlloc(TopMemoryContext, sizeof(PendingRelDelete));

diff --git a/src/backend/catalog/toasting.c b/src/backend/catalog/toasting.c
@@ -198,9 +198,9 @@ create_toast_table(Relation rel, Oid toastOid, Oid toastIndexOid,
 	/*
 	 * Create the toast table and its index
 	 */
-	if (sql_dialect == SQL_DIALECT_TSQL && RelationIsBBFTableVariable(rel))
+	if (IsTsqlTableVariable(rel))
 		pg_toast_prefix = "@pg_toast";
-	else if (sql_dialect == SQL_DIALECT_TSQL && rel->rd_rel->relpersistence == RELPERSISTENCE_TEMP && OidBelongsToENRTempTable(rel->rd_id))
+	else if (IsTsqlTempTable(rel->rd_rel->relpersistence) && OidBelongsToENRTempTable(rel->rd_id))
 		pg_toast_prefix = "#pg_toast";
 
 	snprintf(toast_relname, sizeof(toast_relname),

diff --git a/src/backend/commands/cluster.c b/src/backend/commands/cluster.c
@@ -740,7 +740,7 @@ make_new_heap(Oid OIDOldHeap, Oid NewTableSpace, Oid NewAccessMethod,
 	 * We also must ensure that these temp tables are properly named in TSQL
 	 * so that the metadata is properly cleaned up after in this function.
 	 */
-	if (sql_dialect == SQL_DIALECT_TSQL && relpersistence == RELPERSISTENCE_TEMP && OidBelongsToENRTempTable(OIDOldHeap))
+	if (IsTsqlTempTable(relpersistence) && OidBelongsToENRTempTable(OIDOldHeap))
 		snprintf(NewHeapName, sizeof(NewHeapName), "#pg_temp_%u", OIDOldHeap);
 	else
 		snprintf(NewHeapName, sizeof(NewHeapName), "pg_temp_%u", OIDOldHeap);
@@ -1612,9 +1612,9 @@ finish_heap_swap(Oid OIDOldHeap, Oid OIDNewHeap,
 											 NoLock);
 
 			/* rename the toast table ... */
-			if (sql_dialect == SQL_DIALECT_TSQL && RelationIsBBFTableVariable(newrel))
+			if (IsTsqlTableVariable(newrel))
 				pg_toast_prefix = "@pg_toast";
-			else if (sql_dialect == SQL_DIALECT_TSQL && newrel->rd_rel->relpersistence == RELPERSISTENCE_TEMP && OidBelongsToENRTempTable(newrel->rd_id))
+			else if (IsTsqlTempTable(newrel->rd_rel->relpersistence) && OidBelongsToENRTempTable(newrel->rd_id))
 				pg_toast_prefix = "#pg_toast";
 
 			snprintf(NewToastName, NAMEDATALEN, "%s_%u",

diff --git a/src/backend/utils/activity/pgstat_relation.c b/src/backend/utils/activity/pgstat_relation.c
@@ -174,7 +174,7 @@ void
 pgstat_create_relation(Relation rel)
 {
 	/* Skip pg_stat */
-	if (sql_dialect == SQL_DIALECT_TSQL && rel->rd_rel->relpersistence == RELPERSISTENCE_TEMP && temp_oid_buffer_size > 0)
+	if (IsTsqlTempTable(rel->rd_rel->relpersistence) && UseTempOidBuffer())
 		return;
 	pgstat_create_transactional(PGSTAT_KIND_RELATION,
 								rel->rd_rel->relisshared ? InvalidOid : MyDatabaseId,
@@ -190,7 +190,7 @@ pgstat_drop_relation(Relation rel)
 	int			nest_level = GetCurrentTransactionNestLevel();
 	PgStat_TableStatus *pgstat_info;
 
-	if (sql_dialect == SQL_DIALECT_TSQL && rel->rd_rel->relpersistence == RELPERSISTENCE_TEMP && temp_oid_buffer_size > 0)
+	if (IsTsqlTempTable(rel->rd_rel->relpersistence) && UseTempOidBuffer())
 		return;
 	pgstat_drop_transactional(PGSTAT_KIND_RELATION,
 							  rel->rd_rel->relisshared ? InvalidOid : MyDatabaseId,

diff --git a/src/backend/utils/misc/queryenvironment.c b/src/backend/utils/misc/queryenvironment.c
@@ -1635,20 +1635,35 @@ bool SIMessageIsForTempTable(const SharedInvalidationMessage *msg)
 		elog(ERROR, "unrecognized SI message ID: %d", msg->id);
 }
 
+/* Simple wrapper for Table Variable checks */
+bool IsTsqlTableVariable(Relation rel)
+{
+	return sql_dialect == SQL_DIALECT_TSQL 
+		&& rel
+		&& RelationIsBBFTableVariable(rel);
+}
+
+/* Simple wrapper for Temp Table checks */
+bool IsTsqlTempTable(char relpersistence)
+{
+	return sql_dialect == SQL_DIALECT_TSQL 
+		&& relpersistence == RELPERSISTENCE_TEMP;
+}
+
 /* 
  * Simple check for whether to use temp OID buffer.
  */
-bool useTempOidBuffer()
+bool UseTempOidBuffer()
 {
 	return sql_dialect == SQL_DIALECT_TSQL 
 		&& GetNewTempOidWithIndex_hook 
 		&& temp_oid_buffer_size > 0;
 }
 
 /* Simple check for whether to use temp OID buffer given an existing OID. */
-bool useTempOidBufferForOid(Oid relId)
+bool UseTempOidBufferForOid(Oid relId)
 {
-	return useTempOidBuffer()
+	return UseTempOidBuffer()
 		&& OidBelongsToENRTempTable(relId);
 }
 

diff --git a/src/include/utils/queryenvironment.h b/src/include/utils/queryenvironment.h
@@ -159,8 +159,10 @@ extern void ClearSavedCatcacheMessages(void);
 extern bool SIMessageIsForTempTable(const SharedInvalidationMessage *msg);
 
 /* Various checks */
-extern bool useTempOidBuffer(void);
-extern bool useTempOidBufferForOid(Oid relId);
+extern bool IsTsqlTableVariable(Relation rel);
+extern bool IsTsqlTempTable(char relpersistence);
+extern bool UseTempOidBuffer(void);
+extern bool UseTempOidBufferForOid(Oid relId);
 extern bool has_existing_enr_relations(void);
 
 /* Hooks */
-Original file line number
+Diff line change
@@ Expand Up @@
     	 */
     	rlocator.backend = backend;
-    	use_bbf_oid_buffer = (relpersistence == RELPERSISTENCE_TEMP && useTempOidBuffer());
+    	use_bbf_oid_buffer = (relpersistence == RELPERSISTENCE_TEMP && UseTempOidBuffer());
     	do
     	{
@@ Expand Down @@