Add pivot struct to improve code readability

Added tsql_pivot_fields stuct to store pivot related information to improve readability Signed-off-by: Yanjie Xu <[email protected]>
babelfish-for-postgresql · Nov 21, 2023 · 4362e2a · 4362e2a
1 parent 4c26d9b
commit 4362e2a
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 27 deletions.
diff --git a/contrib/babelfishpg_tsql/runtime/functions.c b/contrib/babelfishpg_tsql/runtime/functions.c
@@ -169,7 +169,7 @@ int 		SPI_execute_raw_parsetree(RawStmt *parsetree, const char *sourcetext, bool
 static HTAB *load_categories_hash(RawStmt *cats_sql, const char *sourcetext, MemoryContext per_query_ctx);
 static Tuplestorestate *get_bbf_pivot_tuplestore(RawStmt 	*sql,
 												const char 	*sourcetext,
-												String 		*funcName,
+												const char 	*funcName,
 												HTAB 		*bbf_pivot_hash,
 												TupleDesc 	tupdesc,
 												bool 		randomAccess);
@@ -3807,15 +3807,14 @@ bbf_pivot(PG_FUNCTION_ARGS)
 	MemoryContext 	per_query_ctx;
 	MemoryContext 	oldcontext;
 	HTAB	   	   	*bbf_pivot_hash;
-
-	MemoryContext 	tsql_outmost_context;
+	MemoryContext 		tsql_outmost_context;
 	PLtsql_execstate 	*tsql_outmost_estat;
-	RawStmt	   		*bbf_pivot_src_sql;
-	RawStmt	   		*bbf_pivot_cat_sql;
-	int				nestlevel;
-	List			*per_pivot_list;
-	char			*query_string;
-	String			*funcName;
+	RawStmt	   			*bbf_pivot_src_sql;
+	RawStmt	   			*bbf_pivot_cat_sql;
+	int					nestlevel;
+	tsql_pivot_fields	*per_pivot_fields;
+	char				*query_string;
+	char				*funcName;
 
 	/* check to see if caller supports us returning a tuplestore */
 	if (rsinfo == NULL || !IsA(rsinfo, ReturnSetInfo))
@@ -3843,12 +3842,11 @@ bbf_pivot(PG_FUNCTION_ARGS)
 	PG_TRY();
 	{
 		Assert(tsql_outmost_estat->pivot_parsetree_list && list_length(tsql_outmost_estat->pivot_parsetree_list) > 0);
-		per_pivot_list = list_nth_node(List, tsql_outmost_estat->pivot_parsetree_list, 0);
-		Assert(list_length(per_pivot_list) >= 2);
-		bbf_pivot_src_sql = list_nth_node(RawStmt, per_pivot_list, 0);
-		bbf_pivot_cat_sql = list_nth_node(RawStmt, per_pivot_list, 1);
-		query_string = list_nth(per_pivot_list, 2);
-		funcName = list_nth(per_pivot_list, 3);
+		per_pivot_fields = (tsql_pivot_fields *) list_nth(tsql_outmost_estat->pivot_parsetree_list, 0);
+		bbf_pivot_src_sql = per_pivot_fields->s_sql;
+		bbf_pivot_cat_sql = per_pivot_fields->c_sql;
+		query_string = per_pivot_fields->sourcetext;
+		funcName = per_pivot_fields->funcName;
 	}
 	PG_CATCH();
 	{
@@ -4005,7 +4003,7 @@ load_categories_hash(RawStmt *cats_sql, const char * sourcetext, MemoryContext p
 static Tuplestorestate *
 get_bbf_pivot_tuplestore(RawStmt 	*sql,
 						const char 	*sourcetext,
-						String		*funcName,
+						const char	*funcName,
 						HTAB 		*bbf_pivot_hash,
 						TupleDesc 	tupdesc,
 						bool 		randomAccess)
@@ -4044,7 +4042,7 @@ get_bbf_pivot_tuplestore(RawStmt 	*sql,
 		int			non_pivot_columns;
 		int			result_ncols;
 		/* only COUNT will output 0 when no row is selected */
-		bool		output_zero = !strcmp(funcName->sval, "count");
+		bool		output_zero = !strcmp(funcName, "count");
 
 		if (num_categories == 0)
 		{

diff --git a/contrib/babelfishpg_tsql/src/hooks.c b/contrib/babelfishpg_tsql/src/hooks.c
@@ -4309,13 +4309,13 @@ transform_pivot_clause(ParseState *pstate, SelectStmt *stmt)
 	TargetEntry	*aggfunc_te;
 	RangeFunction	*pivot_from_function;
 	SelectStmt 		*pivot_src_sql;
-	RawStmt	*s_sql;
-	RawStmt	*c_sql;
-	MemoryContext 	oldContext;
-	MemoryContext 	tsql_outmost_context;
+	RawStmt			*s_sql;
+	RawStmt			*c_sql;
+	tsql_pivot_fields	*pivot_fields;
+	MemoryContext 		oldContext;
+	MemoryContext 		tsql_outmost_context;
 	PLtsql_execstate 	*tsql_outmost_estat;
-	int				nestlevel;
-	char			*sourcetext;
+	int					nestlevel;
 
 	if (sql_dialect != SQL_DIALECT_TSQL)
 		return;
@@ -4450,8 +4450,6 @@ transform_pivot_clause(ParseState *pstate, SelectStmt *stmt)
 			(errcode(ERRCODE_SYNTAX_ERROR),
 				errmsg("pivot outer context not found")));
 
-	oldContext = MemoryContextSwitchTo(tsql_outmost_context);
-	/* save rewrited sqls to global variable for later retrive */
 	s_sql = makeNode(RawStmt);
 	c_sql = makeNode(RawStmt);
 	s_sql->stmt = (Node *) pivot_src_sql;
@@ -4461,9 +4459,15 @@ transform_pivot_clause(ParseState *pstate, SelectStmt *stmt)
 	c_sql->stmt = (Node *) stmt->catSql;
 	c_sql->stmt_location = 0;
 	c_sql->stmt_len = 0;
-	sourcetext = pstrdup(pstate->p_sourcetext);
 
-	tsql_outmost_estat->pivot_parsetree_list = lappend(tsql_outmost_estat->pivot_parsetree_list, list_make4(copyObject(s_sql), copyObject(c_sql), sourcetext, copyObject(funcName)));
+	oldContext = MemoryContextSwitchTo(tsql_outmost_context);
+	/* save rewrited sqls to global variable for later retrive */
+	pivot_fields = (tsql_pivot_fields *) palloc(sizeof(tsql_pivot_fields));
+	pivot_fields->s_sql = copyObject(s_sql);
+	pivot_fields->c_sql = copyObject(c_sql);
+	pivot_fields->sourcetext = pstrdup(pstate->p_sourcetext);
+	pivot_fields->funcName = pstrdup(funcName->sval);
+	tsql_outmost_estat->pivot_parsetree_list = lappend(tsql_outmost_estat->pivot_parsetree_list, pivot_fields);
 	tsql_outmost_estat->pivot_number++;	
 	MemoryContextSwitchTo(oldContext);
 }

diff --git a/contrib/babelfishpg_tsql/src/pltsql.h b/contrib/babelfishpg_tsql/src/pltsql.h
@@ -1816,6 +1816,14 @@ typedef struct tsql_identity_insert_fields
 	Oid			schema_oid;
 } tsql_identity_insert_fields;
 
+typedef struct tsql_pivot_fields
+{
+	RawStmt	*s_sql;
+	RawStmt	*c_sql;
+	char 	*sourcetext;
+	char	*funcName;
+}tsql_pivot_fields;
+
 extern tsql_identity_insert_fields tsql_identity_insert;
 extern check_lang_as_clause_hook_type check_lang_as_clause_hook;
 extern write_stored_proc_probin_hook_type write_stored_proc_probin_hook;